1 files changed, 151 insertions, 0 deletions
diff --git a/tests/mlia/test_devices_ethosu_data_collection.py b/tests/mlia/test_devices_ethosu_data_collection.py
new file mode 100644
index 0000000..897cf41
--- /dev/null
+++ b/tests/mlia/test_devices_ethosu_data_collection.py
@@ -0,0 +1,151 @@
+# SPDX-FileCopyrightText: Copyright 2022, Arm Limited and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for the data collection module for Ethos-U."""
+from pathlib import Path
+from unittest.mock import MagicMock
+
+import pytest
+
+from mlia.core.context import Context
+from mlia.core.data_collection import DataCollector
+from mlia.core.errors import FunctionalityNotSupportedError
+from mlia.devices.ethosu.config import EthosUConfiguration
+from mlia.devices.ethosu.data_collection import EthosUOperatorCompatibility
+from mlia.devices.ethosu.data_collection import EthosUOptimizationPerformance
+from mlia.devices.ethosu.data_collection import EthosUPerformance
+from mlia.devices.ethosu.performance import MemoryUsage
+from mlia.devices.ethosu.performance import NPUCycles
+from mlia.devices.ethosu.performance import OptimizationPerformanceMetrics
+from mlia.devices.ethosu.performance import PerformanceMetrics
+from mlia.nn.tensorflow.optimizations.select import OptimizationSettings
+from mlia.tools.vela_wrapper import Operators
+
+
+@pytest.mark.parametrize(
+    "collector, expected_name",
+    [
+        (
+            EthosUOperatorCompatibility,
+            "ethos_u_operator_compatibility",
+        ),
+        (
+            EthosUPerformance,
+            "ethos_u_performance",
+        ),
+        (
+            EthosUOptimizationPerformance,
+            "ethos_u_model_optimizations",
+        ),
+    ],
+)
+def test_collectors_metadata(
+    collector: DataCollector,
+    expected_name: str,
+) -> None:
+    """Test collectors metadata."""
+    assert collector.name() == expected_name
+
+
+def test_operator_compatibility_collector(
+    dummy_context: Context, test_tflite_model: Path
+) -> None:
+    """Test operator compatibility data collector."""
+    device = EthosUConfiguration("ethos-u55-256")
+
+    collector = EthosUOperatorCompatibility(test_tflite_model, device)
+    collector.set_context(dummy_context)
+
+    result = collector.collect_data()
+    assert isinstance(result, Operators)
+
+
+def test_performance_collector(
+    monkeypatch: pytest.MonkeyPatch, dummy_context: Context, test_tflite_model: Path
+) -> None:
+    """Test performance data collector."""
+    device = EthosUConfiguration("ethos-u55-256")
+
+    mock_performance_estimation(monkeypatch, device)
+
+    collector = EthosUPerformance(test_tflite_model, device)
+    collector.set_context(dummy_context)
+
+    result = collector.collect_data()
+    assert isinstance(result, PerformanceMetrics)
+
+
+def test_optimization_performance_collector(
+    monkeypatch: pytest.MonkeyPatch,
+    dummy_context: Context,
+    test_keras_model: Path,
+    test_tflite_model: Path,
+) -> None:
+    """Test optimization performance data collector."""
+    device = EthosUConfiguration("ethos-u55-256")
+
+    mock_performance_estimation(monkeypatch, device)
+    collector = EthosUOptimizationPerformance(
+        test_keras_model,
+        device,
+        [
+            [
+                {"optimization_type": "pruning", "optimization_target": 0.5},
+            ]
+        ],
+    )
+    collector.set_context(dummy_context)
+    result = collector.collect_data()
+
+    assert isinstance(result, OptimizationPerformanceMetrics)
+    assert isinstance(result.original_perf_metrics, PerformanceMetrics)
+    assert isinstance(result.optimizations_perf_metrics, list)
+    assert len(result.optimizations_perf_metrics) == 1
+
+    opt, metrics = result.optimizations_perf_metrics[0]
+    assert opt == [OptimizationSettings("pruning", 0.5, None)]
+    assert isinstance(metrics, PerformanceMetrics)
+
+    collector_no_optimizations = EthosUOptimizationPerformance(
+        test_keras_model,
+        device,
+        [],
+    )
+    with pytest.raises(FunctionalityNotSupportedError):
+        collector_no_optimizations.collect_data()
+
+    collector_tflite = EthosUOptimizationPerformance(
+        test_tflite_model,
+        device,
+        [
+            [
+                {"optimization_type": "pruning", "optimization_target": 0.5},
+            ]
+        ],
+    )
+    collector_tflite.set_context(dummy_context)
+    with pytest.raises(FunctionalityNotSupportedError):
+        collector_tflite.collect_data()
+
+    with pytest.raises(
+        Exception, match="Optimization parameters expected to be a list"
+    ):
+        collector_bad_config = EthosUOptimizationPerformance(
+            test_keras_model, device, {"optimization_type": "pruning"}  # type: ignore
+        )
+        collector.set_context(dummy_context)
+        collector_bad_config.collect_data()
+
+
+def mock_performance_estimation(
+    monkeypatch: pytest.MonkeyPatch, device: EthosUConfiguration
+) -> None:
+    """Mock performance estimation."""
+    metrics = PerformanceMetrics(
+        device,
+        NPUCycles(1, 2, 3, 4, 5, 6),
+        MemoryUsage(1, 2, 3, 4, 5),
+    )
+    monkeypatch.setattr(
+        "mlia.devices.ethosu.data_collection.EthosUPerformanceEstimator.estimate",
+        MagicMock(return_value=metrics),
+    )