haoline 0.3.0__py3-none-any.whl

haoline/tests/test_tensorflow.py

@@ -0,0 +1,380 @@
+# Copyright (c) 2025 HaoLine Contributors
+# SPDX-License-Identifier: MIT
+
+"""
+Unit tests for TensorFlow/Keras/JAX to ONNX conversion functionality.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import logging
+from unittest.mock import patch
+
+import pytest
+
+from ..cli import (
+    _convert_frozen_graph_to_onnx,
+    _convert_jax_to_onnx,
+    _convert_keras_to_onnx,
+    _convert_tensorflow_to_onnx,
+)
+
+# Check if TensorFlow and tf2onnx are available
+_TF_AVAILABLE = (
+    importlib.util.find_spec("tensorflow") is not None
+    and importlib.util.find_spec("tf2onnx") is not None
+)
+_JAX_AVAILABLE = importlib.util.find_spec("jax") is not None
+
+# Import TensorFlow only when needed for tests
+if _TF_AVAILABLE:
+    import tensorflow as tf
+
+
+@pytest.fixture
+def logger():
+    """Create a logger for tests."""
+    return logging.getLogger("test_tensorflow")
+
+
+class TestTensorFlowConversionErrors:
+    """Test error handling for TensorFlow conversion (no TF required)."""
+
+    def test_nonexistent_savedmodel(self, tmp_path, logger):
+        """Conversion should fail for non-existent SavedModel."""
+        fake_path = tmp_path / "nonexistent_model"
+
+        onnx_path, _ = _convert_tensorflow_to_onnx(
+            fake_path,
+            output_path=None,
+            opset_version=17,
+            logger=logger,
+        )
+
+        assert onnx_path is None
+
+    def test_tf2onnx_not_installed(self, tmp_path, logger):
+        """Conversion should fail gracefully when tf2onnx not installed."""
+        # Create a fake directory to avoid "not found" error
+        fake_model = tmp_path / "fake_saved_model"
+        fake_model.mkdir()
+
+        # Mock tf2onnx import to fail
+        with patch.dict("sys.modules", {"tf2onnx": None}):
+            # Force re-import check by clearing cached imports
+
+            # Save original function
+
+            # The function checks for import at runtime, so we need to
+            # test the actual error path
+            _onnx_path, _ = _convert_tensorflow_to_onnx(
+                fake_model,
+                output_path=None,
+                opset_version=17,
+                logger=logger,
+            )
+
+            # If tf2onnx is not installed, this should fail
+            # (result depends on actual installation state)
+
+
+class TestKerasConversionErrors:
+    """Test error handling for Keras conversion (no TF required)."""
+
+    def test_nonexistent_keras_file(self, tmp_path, logger):
+        """Conversion should fail for non-existent Keras file."""
+        fake_path = tmp_path / "nonexistent.h5"
+
+        onnx_path, _ = _convert_keras_to_onnx(
+            fake_path,
+            output_path=None,
+            opset_version=17,
+            logger=logger,
+        )
+
+        assert onnx_path is None
+
+    def test_unexpected_extension_warning(self, tmp_path, logger, caplog):
+        """Unexpected file extension should log warning but proceed."""
+        # Create a fake file with wrong extension
+        fake_file = tmp_path / "model.wrongext"
+        fake_file.write_text("fake content")
+
+        with caplog.at_level(logging.WARNING):
+            _onnx_path, _ = _convert_keras_to_onnx(
+                fake_file,
+                output_path=None,
+                opset_version=17,
+                logger=logger,
+            )
+
+        # Should have logged a warning about unexpected extension
+        # (will still fail because it's not a real model)
+
+
+class TestFrozenGraphConversionErrors:
+    """Test error handling for frozen graph conversion."""
+
+    def test_missing_inputs_outputs(self, tmp_path, logger):
+        """Conversion should fail when inputs/outputs not specified."""
+        fake_pb = tmp_path / "model.pb"
+        fake_pb.write_bytes(b"fake")
+
+        # Missing inputs
+        onnx_path, _ = _convert_frozen_graph_to_onnx(
+            fake_pb,
+            inputs=None,
+            outputs="output:0",
+            output_path=None,
+            opset_version=17,
+            logger=logger,
+        )
+        assert onnx_path is None
+
+        # Missing outputs
+        onnx_path, _ = _convert_frozen_graph_to_onnx(
+            fake_pb,
+            inputs="input:0",
+            outputs=None,
+            output_path=None,
+            opset_version=17,
+            logger=logger,
+        )
+        assert onnx_path is None
+
+    def test_nonexistent_pb_file(self, tmp_path, logger):
+        """Conversion should fail for non-existent .pb file."""
+        fake_path = tmp_path / "nonexistent.pb"
+
+        onnx_path, _ = _convert_frozen_graph_to_onnx(
+            fake_path,
+            inputs="input:0",
+            outputs="output:0",
+            output_path=None,
+            opset_version=17,
+            logger=logger,
+        )
+
+        assert onnx_path is None
+
+
+class TestJAXConversionErrors:
+    """Test error handling for JAX conversion."""
+
+    def test_missing_apply_fn(self, tmp_path, logger):
+        """Conversion should fail when --jax-apply-fn not provided."""
+        fake_params = tmp_path / "params.pkl"
+        fake_params.write_bytes(b"fake")
+
+        onnx_path, _ = _convert_jax_to_onnx(
+            fake_params,
+            apply_fn_path=None,  # Missing
+            input_shape_str="1,3,224,224",
+            output_path=None,
+            opset_version=17,
+            logger=logger,
+        )
+
+        assert onnx_path is None
+
+    def test_missing_input_shape(self, tmp_path, logger):
+        """Conversion should fail when --input-shape not provided."""
+        fake_params = tmp_path / "params.pkl"
+        fake_params.write_bytes(b"fake")
+
+        onnx_path, _ = _convert_jax_to_onnx(
+            fake_params,
+            apply_fn_path="module:function",
+            input_shape_str=None,  # Missing
+            output_path=None,
+            opset_version=17,
+            logger=logger,
+        )
+
+        assert onnx_path is None
+
+    def test_invalid_apply_fn_format(self, tmp_path, logger):
+        """Conversion should fail for invalid apply_fn format."""
+        fake_params = tmp_path / "params.pkl"
+        fake_params.write_bytes(b"fake")
+
+        onnx_path, _ = _convert_jax_to_onnx(
+            fake_params,
+            apply_fn_path="no_colon_separator",  # Invalid format
+            input_shape_str="1,3,224,224",
+            output_path=None,
+            opset_version=17,
+            logger=logger,
+        )
+
+        assert onnx_path is None
+
+    def test_invalid_input_shape_format(self, tmp_path, logger):
+        """Conversion should fail for invalid input shape format."""
+        fake_params = tmp_path / "params.pkl"
+        fake_params.write_bytes(b"fake")
+
+        onnx_path, _ = _convert_jax_to_onnx(
+            fake_params,
+            apply_fn_path="module:function",
+            input_shape_str="not,valid,shape",  # Invalid
+            output_path=None,
+            opset_version=17,
+            logger=logger,
+        )
+
+        assert onnx_path is None
+
+    def test_nonexistent_params_file(self, tmp_path, logger):
+        """Conversion should fail for non-existent params file."""
+        fake_path = tmp_path / "nonexistent.pkl"
+
+        onnx_path, _ = _convert_jax_to_onnx(
+            fake_path,
+            apply_fn_path="module:function",
+            input_shape_str="1,3,224,224",
+            output_path=None,
+            opset_version=17,
+            logger=logger,
+        )
+
+        assert onnx_path is None
+
+
+@pytest.mark.skipif(not _TF_AVAILABLE, reason="TensorFlow not installed")
+class TestTensorFlowConversion:
+    """Tests for TensorFlow to ONNX conversion (requires TensorFlow)."""
+
+    def test_savedmodel_conversion(self, tmp_path, logger):
+        """SavedModel should convert successfully."""
+        # Create a simple TF SavedModel
+        model = tf.keras.Sequential(
+            [
+                tf.keras.layers.Dense(10, input_shape=(5,), activation="relu"),
+                tf.keras.layers.Dense(2),
+            ]
+        )
+
+        saved_model_path = tmp_path / "saved_model"
+        model.save(str(saved_model_path), save_format="tf")
+
+        # Convert to ONNX
+        onnx_path, _temp_file = _convert_tensorflow_to_onnx(
+            saved_model_path,
+            output_path=tmp_path / "output.onnx",
+            opset_version=17,
+            logger=logger,
+        )
+
+        assert onnx_path is not None
+        assert onnx_path.exists()
+        assert onnx_path.suffix == ".onnx"
+
+    def test_savedmodel_temp_file(self, tmp_path, logger):
+        """SavedModel conversion to temp file should work."""
+        model = tf.keras.Sequential(
+            [
+                tf.keras.layers.Dense(10, input_shape=(5,)),
+            ]
+        )
+
+        saved_model_path = tmp_path / "saved_model"
+        model.save(str(saved_model_path), save_format="tf")
+
+        onnx_path, _temp_file = _convert_tensorflow_to_onnx(
+            saved_model_path,
+            output_path=None,  # Use temp file
+            opset_version=17,
+            logger=logger,
+        )
+
+        assert onnx_path is not None
+        assert onnx_path.exists()
+        # Cleanup
+        onnx_path.unlink()
+
+
+@pytest.mark.skipif(not _TF_AVAILABLE, reason="TensorFlow not installed")
+class TestKerasConversion:
+    """Tests for Keras to ONNX conversion (requires TensorFlow)."""
+
+    def test_h5_conversion(self, tmp_path, logger):
+        """Keras .h5 model should convert successfully."""
+        model = tf.keras.Sequential(
+            [
+                tf.keras.layers.Dense(10, input_shape=(5,), activation="relu"),
+                tf.keras.layers.Dense(2),
+            ]
+        )
+
+        h5_path = tmp_path / "model.h5"
+        model.save(str(h5_path), save_format="h5")
+
+        onnx_path, _temp_file = _convert_keras_to_onnx(
+            h5_path,
+            output_path=tmp_path / "output.onnx",
+            opset_version=17,
+            logger=logger,
+        )
+
+        assert onnx_path is not None
+        assert onnx_path.exists()
+        assert onnx_path.suffix == ".onnx"
+
+    def test_keras_format_conversion(self, tmp_path, logger):
+        """Keras .keras format should convert successfully."""
+        model = tf.keras.Sequential(
+            [
+                tf.keras.layers.Dense(10, input_shape=(5,)),
+            ]
+        )
+
+        keras_path = tmp_path / "model.keras"
+        model.save(str(keras_path))
+
+        onnx_path, _temp_file = _convert_keras_to_onnx(
+            keras_path,
+            output_path=tmp_path / "output.onnx",
+            opset_version=17,
+            logger=logger,
+        )
+
+        assert onnx_path is not None
+        assert onnx_path.exists()
+
+
+@pytest.mark.skipif(not _JAX_AVAILABLE or not _TF_AVAILABLE, reason="JAX or TF not installed")
+class TestJAXConversion:
+    """Tests for JAX to ONNX conversion (requires JAX and TensorFlow)."""
+
+    def test_unsupported_params_format(self, tmp_path, logger):
+        """Unsupported params format should fail with clear error."""
+        fake_params = tmp_path / "params.xyz"
+        fake_params.write_bytes(b"fake")
+
+        onnx_path, _ = _convert_jax_to_onnx(
+            fake_params,
+            apply_fn_path="module:function",
+            input_shape_str="1,3,224,224",
+            output_path=None,
+            opset_version=17,
+            logger=logger,
+        )
+
+        assert onnx_path is None
+
+
+class TestCLIValidation:
+    """Test CLI argument validation."""
+
+    def test_multiple_conversion_flags_error(self):
+        """Using multiple conversion flags should error."""
+
+        # This would be tested via the main function, but we can verify
+        # the argument structure allows these to be set
+        # The actual validation happens in run_inspect()
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

haoline/tests/test_visualizations.py

@@ -0,0 +1,316 @@
+# Copyright (c) 2025 HaoLine Contributors
+# SPDX-License-Identifier: MIT
+
+"""
+Unit tests for the visualization module.
+
+Tests chart generation, theming, and graceful fallback behavior.
+"""
+
+from __future__ import annotations
+
+import tempfile
+from pathlib import Path
+
+import pytest
+
+from ..analyzer import FlopCounts, MemoryEstimates, ParamCounts
+from ..report import GraphSummary, InspectionReport, ModelMetadata
+from ..visualizations import (
+    THEME,
+    ChartTheme,
+    VisualizationGenerator,
+    _format_count,
+    generate_visualizations,
+    is_available,
+)
+
+
+class TestVisualizationAvailability:
+    """Tests for matplotlib availability detection."""
+
+    def test_is_available_returns_bool(self):
+        """is_available() should return a boolean."""
+        result = is_available()
+        assert isinstance(result, bool)
+
+    def test_theme_has_required_colors(self):
+        """Theme should have all required color properties."""
+        assert hasattr(THEME, "background")
+        assert hasattr(THEME, "text")
+        assert hasattr(THEME, "accent_primary")
+        assert hasattr(THEME, "palette")
+        assert len(THEME.palette) >= 5  # Need enough colors for charts
+
+
+class TestChartTheme:
+    """Tests for the ChartTheme dataclass."""
+
+    def test_default_theme_values(self):
+        """Default theme should have sensible values."""
+        theme = ChartTheme()
+        assert theme.background.startswith("#")
+        assert theme.text.startswith("#")
+        assert theme.figure_dpi >= 72
+        assert theme.figure_width > 0
+        assert theme.figure_height > 0
+
+    def test_custom_theme(self):
+        """Should be able to create custom themes."""
+        theme = ChartTheme(
+            background="#000000",
+            text="#ffffff",
+            accent_primary="#ff0000",
+        )
+        assert theme.background == "#000000"
+        assert theme.text == "#ffffff"
+        assert theme.accent_primary == "#ff0000"
+
+
+class TestFormatCount:
+    """Tests for the _format_count helper function."""
+
+    def test_format_small_numbers(self):
+        """Small numbers should be formatted as-is."""
+        assert _format_count(0) == "0"
+        assert _format_count(1) == "1"
+        assert _format_count(999) == "999"
+
+    def test_format_thousands(self):
+        """Thousands should use K suffix."""
+        assert _format_count(1000) == "1.0K"
+        assert _format_count(1500) == "1.5K"
+        assert _format_count(999999) == "1000.0K"  # Just under 1M
+
+    def test_format_millions(self):
+        """Millions should use M suffix."""
+        assert _format_count(1000000) == "1.0M"
+        assert _format_count(1500000) == "1.5M"
+        assert _format_count(25000000) == "25.0M"
+
+    def test_format_billions(self):
+        """Billions should use B suffix."""
+        assert _format_count(1000000000) == "1.0B"
+        assert _format_count(7500000000) == "7.5B"
+
+
+class TestVisualizationGenerator:
+    """Tests for the VisualizationGenerator class."""
+
+    def test_generator_initialization(self):
+        """Generator should initialize without errors."""
+        gen = VisualizationGenerator()
+        assert gen is not None
+        assert gen.logger is not None
+
+    def test_generator_with_custom_logger(self):
+        """Generator should accept custom logger."""
+        import logging
+
+        logger = logging.getLogger("test")
+        gen = VisualizationGenerator(logger=logger)
+        assert gen.logger is logger
+
+
+@pytest.mark.skipif(not is_available(), reason="matplotlib not installed")
+class TestChartGeneration:
+    """Tests for actual chart generation (requires matplotlib)."""
+
+    def test_operator_histogram_generation(self):
+        """Should generate operator histogram PNG."""
+        gen = VisualizationGenerator()
+        op_counts = {
+            "Conv": 50,
+            "Relu": 48,
+            "Add": 25,
+            "MatMul": 12,
+            "Softmax": 6,
+        }
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            output_path = Path(tmpdir) / "op_histogram.png"
+            result = gen.operator_histogram(op_counts, output_path)
+
+            assert result is not None
+            assert result.exists()
+            assert result.stat().st_size > 0  # File has content
+
+    def test_param_distribution_generation(self):
+        """Should generate parameter distribution chart."""
+        gen = VisualizationGenerator()
+        params_by_op = {
+            "Conv": 5000000,
+            "Gemm": 2000000,
+            "MatMul": 1500000,
+        }
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            output_path = Path(tmpdir) / "param_dist.png"
+            result = gen.param_distribution(params_by_op, output_path)
+
+            assert result is not None
+            assert result.exists()
+
+    def test_flops_distribution_generation(self):
+        """Should generate FLOPs distribution chart."""
+        gen = VisualizationGenerator()
+        flops_by_op = {
+            "Conv": 500000000,
+            "MatMul": 200000000,
+            "Gemm": 150000000,
+        }
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            output_path = Path(tmpdir) / "flops_dist.png"
+            result = gen.flops_distribution(flops_by_op, output_path)
+
+            assert result is not None
+            assert result.exists()
+
+    def test_empty_data_returns_none(self):
+        """Charts should return None for empty data."""
+        gen = VisualizationGenerator()
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            # Empty dict
+            result = gen.operator_histogram({}, Path(tmpdir) / "empty.png")
+            assert result is None
+
+            # All zeros
+            result = gen.param_distribution({"Conv": 0, "Relu": 0}, Path(tmpdir) / "zeros.png")
+            assert result is None
+
+    def test_generate_all_creates_multiple_charts(self):
+        """generate_all should create multiple chart files."""
+        # Create a mock report
+        metadata = ModelMetadata(
+            path="test.onnx",
+            ir_version=8,
+            producer_name="test",
+            producer_version="1.0",
+            domain="",
+            model_version=1,
+            doc_string="",
+            opsets={"ai.onnx": 17},
+        )
+
+        graph_summary = GraphSummary(
+            num_nodes=100,
+            num_inputs=1,
+            num_outputs=1,
+            num_initializers=50,
+            input_shapes={"input": [1, 3, 224, 224]},
+            output_shapes={"output": [1, 1000]},
+            op_type_counts={"Conv": 50, "Relu": 48, "Add": 25},
+        )
+
+        param_counts = ParamCounts(
+            total=25000000,
+            trainable=25000000,
+            by_op_type={"Conv": 20000000, "Gemm": 5000000},
+        )
+
+        flop_counts = FlopCounts(
+            total=4000000000,
+            by_op_type={"Conv": 3500000000, "Gemm": 500000000},
+        )
+
+        memory_estimates = MemoryEstimates(
+            model_size_bytes=100000000,
+            peak_activation_bytes=50000000,
+        )
+
+        report = InspectionReport(
+            metadata=metadata,
+            graph_summary=graph_summary,
+            param_counts=param_counts,
+            flop_counts=flop_counts,
+            memory_estimates=memory_estimates,
+        )
+
+        gen = VisualizationGenerator()
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            paths = gen.generate_all(report, Path(tmpdir))
+
+            assert len(paths) >= 3  # Should generate at least 3 charts
+            assert "op_histogram" in paths
+            assert "param_distribution" in paths
+            assert "complexity_summary" in paths
+
+            # All files should exist
+            for name, path in paths.items():
+                assert path.exists(), f"{name} file should exist"
+
+
+@pytest.mark.skipif(not is_available(), reason="matplotlib not installed")
+class TestConvenienceFunction:
+    """Tests for the generate_visualizations convenience function."""
+
+    def test_generate_visualizations_function(self):
+        """Convenience function should work correctly."""
+        metadata = ModelMetadata(
+            path="test.onnx",
+            ir_version=8,
+            producer_name="test",
+            producer_version="1.0",
+            domain="",
+            model_version=1,
+            doc_string="",
+            opsets={"ai.onnx": 17},
+        )
+
+        graph_summary = GraphSummary(
+            num_nodes=10,
+            num_inputs=1,
+            num_outputs=1,
+            num_initializers=5,
+            input_shapes={},
+            output_shapes={},
+            op_type_counts={"Conv": 5, "Relu": 5},
+        )
+
+        param_counts = ParamCounts(total=1000, trainable=1000, by_op_type={"Conv": 1000})
+        flop_counts = FlopCounts(total=10000, by_op_type={"Conv": 10000})
+        memory_estimates = MemoryEstimates(model_size_bytes=4000, peak_activation_bytes=2000)
+
+        report = InspectionReport(
+            metadata=metadata,
+            graph_summary=graph_summary,
+            param_counts=param_counts,
+            flop_counts=flop_counts,
+            memory_estimates=memory_estimates,
+        )
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            paths = generate_visualizations(report, tmpdir)
+            assert isinstance(paths, dict)
+            assert len(paths) > 0
+
+
+class TestGracefulDegradation:
+    """Tests for graceful degradation when matplotlib is unavailable."""
+
+    def test_generator_handles_missing_matplotlib(self):
+        """Generator should not crash if matplotlib is unavailable."""
+        # This test runs regardless of matplotlib availability
+        gen = VisualizationGenerator()
+
+        # Even with matplotlib, empty data should return empty dict
+        metadata = ModelMetadata(
+            path="test.onnx",
+            ir_version=8,
+            producer_name="test",
+            producer_version="1.0",
+            domain="",
+            model_version=1,
+            doc_string="",
+            opsets={},
+        )
+
+        # Report with no metrics - should return empty dict
+        report = InspectionReport(metadata=metadata)
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            paths = gen.generate_all(report, Path(tmpdir))
+            assert isinstance(paths, dict)