aboutsummaryrefslogtreecommitdiff
path: root/tests/test_cli_helpers.py
blob: 494ed8958b3b6c26f3f43bbbe6e00b2727eda3ab (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
# SPDX-FileCopyrightText: Copyright 2022-2023, Arm Limited and/or its affiliates.
# SPDX-License-Identifier: Apache-2.0
"""Tests for the helper classes."""
from __future__ import annotations

from pathlib import Path
from typing import Any

import pytest

from mlia.cli.helpers import CLIActionResolver
from mlia.cli.helpers import copy_profile_file_to_output_dir
from mlia.nn.select import OptimizationSettings


class TestCliActionResolver:
    """Test cli action resolver."""

    @staticmethod
    @pytest.mark.parametrize(
        "args, params, expected_result",
        [
            [
                {},
                {"opt_settings": "some_setting"},
                [],
            ],
            [
                {},
                {},
                [
                    "Note: you will need a Keras model for that.",
                    "For example: mlia optimize /path/to/keras_model "
                    "--pruning --clustering "
                    "--pruning-target 0.5 --clustering-target 32",
                    "For more info: mlia optimize --help",
                ],
            ],
            [
                {"model": "model.h5"},
                {},
                [
                    "For example: mlia optimize model.h5 --pruning --clustering "
                    "--pruning-target 0.5 --clustering-target 32",
                    "For more info: mlia optimize --help",
                ],
            ],
            [
                {"model": "model.h5"},
                {"opt_settings": [OptimizationSettings("pruning", 0.5, None)]},
                [
                    "For more info: mlia optimize --help",
                    "Optimization command: "
                    "mlia optimize model.h5 --pruning "
                    "--pruning-target 0.5",
                ],
            ],
            [
                {"model": "model.h5", "target_profile": "target_profile"},
                {"opt_settings": [OptimizationSettings("pruning", 0.5, None)]},
                [
                    "For more info: mlia optimize --help",
                    "Optimization command: "
                    "mlia optimize model.h5 --target-profile target_profile "
                    "--pruning --pruning-target 0.5",
                ],
            ],
        ],
    )
    def test_apply_optimizations(
        args: dict[str, Any],
        params: dict[str, Any],
        expected_result: list[str],
    ) -> None:
        """Test action resolving for applying optimizations."""
        resolver = CLIActionResolver(args)
        assert resolver.apply_optimizations(**params) == expected_result

    @staticmethod
    def test_operator_compatibility_details() -> None:
        """Test operator compatibility details info."""
        resolver = CLIActionResolver({})
        assert resolver.operator_compatibility_details() == [
            "For more details, run: mlia check --help"
        ]

    @staticmethod
    def test_optimization_details() -> None:
        """Test optimization details info."""
        resolver = CLIActionResolver({})
        assert resolver.optimization_details() == [
            "For more info, see: mlia optimize --help"
        ]

    @staticmethod
    @pytest.mark.parametrize(
        "args, expected_result",
        [
            [
                {},
                [],
            ],
            [
                {"model": "model.tflite", "target_profile": "target_profile"},
                [
                    "Check the estimated performance by running the "
                    "following command: ",
                    "mlia check model.tflite "
                    "--target-profile target_profile --performance",
                ],
            ],
        ],
    )
    def test_check_performance(
        args: dict[str, Any], expected_result: list[str]
    ) -> None:
        """Test check performance info."""
        resolver = CLIActionResolver(args)
        assert resolver.check_performance() == expected_result

    @staticmethod
    @pytest.mark.parametrize(
        "args, expected_result",
        [
            [
                {},
                [],
            ],
            [
                {"model": "model.tflite", "target_profile": "target_profile"},
                [
                    "Try running the following command to verify that:",
                    "mlia check model.tflite --target-profile target_profile",
                ],
            ],
        ],
    )
    def test_check_operator_compatibility(
        args: dict[str, Any], expected_result: list[str]
    ) -> None:
        """Test checking operator compatibility info."""
        resolver = CLIActionResolver(args)
        assert resolver.check_operator_compatibility() == expected_result


def test_copy_profile_file_to_output_dir(tmp_path: Path) -> None:
    """Test if the profile file is copied into the output directory."""
    test_target_profile_name = "ethos-u55-128"
    test_file_path = Path(f"{tmp_path}/{test_target_profile_name}.toml")

    copy_profile_file_to_output_dir(test_target_profile_name, tmp_path)
    assert Path.is_file(test_file_path)