diff --git a/tests/integration/cli/test_cli.py b/tests/integration/cli/test_cli.py index 1385d756c7..e524932946 100644 --- a/tests/integration/cli/test_cli.py +++ b/tests/integration/cli/test_cli.py @@ -12,7 +12,7 @@ import torch from anomalib.cli import AnomalibCLI -from anomalib.deploy import ExportType +from anomalib.deploy import CompressionType, ExportType class TestCLI: @@ -155,11 +155,21 @@ def test_predict_with_image_path(self, project_path: Path) -> None: ) torch.cuda.empty_cache() - @pytest.mark.parametrize("export_type", [ExportType.TORCH, ExportType.ONNX, ExportType.OPENVINO]) + @pytest.mark.parametrize( + ("export_type", "compression_type"), + [ + (ExportType.TORCH, None), + (ExportType.ONNX, None), + (ExportType.OPENVINO, None), + (ExportType.OPENVINO, CompressionType.POT), + (ExportType.OPENVINO, CompressionType.NNCF), + ], + ) def test_export( self, project_path: Path, export_type: ExportType, + compression_type: CompressionType | None, ) -> None: """Test the export method of the CLI. @@ -167,17 +177,19 @@ def test_export( dataset_path (Path): Root of the synthetic/original dataset. project_path (Path): Path to temporary project folder. export_type (ExportType): Export type. + compression_type (CompressionType): Compression type (if any). """ - AnomalibCLI( - args=[ - "export", - "--export_type", - export_type, - *self._get_common_cli_args(None, project_path), - "--ckpt_path", - f"{project_path}/Padim/MVTec/dummy/v0/weights/lightning/model.ckpt", - ], - ) + args = [ + "export", + "--export_type", + export_type, + *self._get_common_cli_args(None, project_path), + "--ckpt_path", + f"{project_path}/Padim/MVTec/dummy/v0/weights/lightning/model.ckpt", + ] + if compression_type: + args += ["--compression_type", str(compression_type)] + AnomalibCLI(args=args) @staticmethod def _get_common_cli_args(dataset_path: Path | None, project_path: Path) -> list[str]: