feat: add unit tests for CLI, GPUAnalyzer, and human-readable utilities

Author: PythonicVarun

tests/test_cli.py

@@ -0,0 +1,101 @@
+import unittest
+from unittest.mock import MagicMock, patch
+
+from click.testing import CliRunner
+
+from canirun.cli import main
+from canirun.enum import COMPATIBILITY
+
+
+class TestCli(unittest.TestCase):
+    """Test suite for the CLI module."""
+
+    def setUp(self) -> None:
+        """Sets up the test environment."""
+        self.runner = CliRunner()
+
+    @patch("canirun.cli.ModelAnalyzer")
+    def test_cli_success(self, MockAnalyzer: MagicMock) -> None:
+        """Tests the CLI with a successful model analysis."""
+        # Setup mock
+        instance = MockAnalyzer.return_value
+        instance.fetch_model_data.return_value = {"some": "data"}
+        instance.specs = {
+            "ram": 16 * 1024**3,
+            "vram": 8 * 1024**3,
+            "name": "Test GPU",
+            "is_mac": False,
+        }
+        instance.calculate.return_value = [
+            {
+                "quant": "FP16",
+                "total_ram": 14 * 1024**3,
+                "kv_cache": 1 * 1024**3,
+                "status": COMPATIBILITY.FULL,
+            }
+        ]
+
+        # Run CLI
+        result = self.runner.invoke(main, ["test-model"])
+
+        # Verify
+        self.assertEqual(result.exit_code, 0)
+        self.assertIn("ANALYSIS REPORT: test-model", result.output)
+        self.assertIn("FP16", result.output)
+        self.assertIn("✅ GPU", result.output)
+
+    @patch("canirun.cli.ModelAnalyzer")
+    def test_cli_fetch_error(self, MockAnalyzer: MagicMock) -> None:
+        """Tests the CLI when model data fetch fails."""
+        instance = MockAnalyzer.return_value
+        instance.fetch_model_data.return_value = None
+
+        result = self.runner.invoke(main, ["bad-model"])
+
+        self.assertEqual(result.exit_code, 0)
+        self.assertIn("Error: Could not fetch data", result.output)
+
+    @patch("canirun.cli.ModelAnalyzer")
+    def test_cli_verbose(self, MockAnalyzer: MagicMock) -> None:
+        """Tests the CLI with verbose flag."""
+        instance = MockAnalyzer.return_value
+        instance.fetch_model_data.return_value = {"some": "data"}
+        instance.specs = {"ram": 0, "vram": 0, "name": "Test", "is_mac": False}
+        instance.calculate.return_value = [
+            {
+                "quant": "FP16",
+                "total_ram": 100,
+                "kv_cache": 10,
+                "status": COMPATIBILITY.FULL,
+            }
+        ]
+
+        result = self.runner.invoke(main, ["test-model", "--verbose"])
+
+        self.assertEqual(result.exit_code, 0)
+        # Verify ModelAnalyzer was initialized with verbose=True
+        MockAnalyzer.assert_called_with("test-model", verbose=True, hf_token=None)
+
+    @patch("canirun.cli.ModelAnalyzer")
+    def test_cli_with_token(self, MockAnalyzer: MagicMock) -> None:
+        """Tests the CLI with an HF token."""
+        instance = MockAnalyzer.return_value
+        instance.fetch_model_data.return_value = {"some": "data"}
+        instance.specs = {"ram": 0, "vram": 0, "name": "Test", "is_mac": False}
+        instance.calculate.return_value = [
+            {
+                "quant": "FP16",
+                "total_ram": 100,
+                "kv_cache": 10,
+                "status": COMPATIBILITY.FULL,
+            }
+        ]
+
+        result = self.runner.invoke(main, ["test-model", "--hf-token", "hf_123"])
+
+        self.assertEqual(result.exit_code, 0)
+        MockAnalyzer.assert_called_with("test-model", verbose=False, hf_token="hf_123")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_gpu.py

@@ -0,0 +1,138 @@
+import sys
+import unittest
+from unittest.mock import MagicMock, patch
+
+from canirun.gpu import GPUAnalyzer
+
+
+class TestGPUAnalyzer(unittest.TestCase):
+    """Test suite for the GPUAnalyzer class."""
+
+    @patch("canirun.gpu.shutil.which")
+    @patch("subprocess.run")
+    def test_gpu_detection_linux_success(
+        self,
+        mock_run: MagicMock,
+        mock_which: MagicMock,
+    ) -> None:
+        """Tests successful GPU detection on Linux."""
+        # Simulate Linux environment
+        with patch("sys.platform", "linux"):
+            mock_which.return_value = "/usr/bin/nvidia-smi"
+
+            # Mock successful subprocess output
+            mock_run.return_value.stdout = "0, NVIDIA GeForce RTX 3080, 10240"
+            mock_run.return_value.returncode = 0
+
+            analyzer = GPUAnalyzer(verbose=True)
+
+            self.assertTrue(analyzer.is_gpu_available())
+            self.assertEqual(analyzer.gpu_count, 1)
+            self.assertEqual(analyzer.device_name, "NVIDIA GeForce RTX 3080")
+            self.assertEqual(analyzer.get_device_name(0), "NVIDIA GeForce RTX 3080")
+
+            # 10240 MiB -> bytes
+            expected_vram = 10240 * 1024 * 1024
+            self.assertEqual(analyzer.vram, expected_vram)
+            self.assertEqual(analyzer.get_vram(0), expected_vram)
+
+            # Check subprocess call
+            mock_run.assert_called_once()
+            args, _ = mock_run.call_args
+            self.assertEqual(args[0][0], "nvidia-smi")
+
+    @patch("canirun.gpu.shutil.which")
+    @patch("subprocess.run")
+    def test_gpu_detection_windows_success(
+        self,
+        mock_run: MagicMock,
+        mock_which: MagicMock,
+    ) -> None:
+        """Tests successful GPU detection on Windows."""
+        # Simulate Windows environment
+        with patch("sys.platform", "win32"):
+            nvidia_path = "C:\\Program Files\\NVIDIA Corporation\\NVSMI\\nvidia-smi.exe"
+            mock_which.return_value = nvidia_path
+
+            mock_run.return_value.stdout = "0, NVIDIA GeForce RTX 4090, 24576"
+            mock_run.return_value.returncode = 0
+
+            analyzer = GPUAnalyzer(verbose=False)
+
+            self.assertTrue(analyzer.is_gpu_available())
+            self.assertEqual(analyzer.device_name, "NVIDIA GeForce RTX 4090")
+
+            # 24576 MiB -> bytes
+            expected_vram = 24576 * 1024 * 1024
+            self.assertEqual(analyzer.vram, expected_vram)
+            self.assertEqual(analyzer.get_vram(0), expected_vram)
+
+            # Check subprocess call used full path or found executable
+            mock_run.assert_called_once()
+            args, _ = mock_run.call_args
+            self.assertEqual(args[0][0], nvidia_path)
+
+    @patch("canirun.gpu.shutil.which")
+    @patch("subprocess.run")
+    def test_gpu_not_found(self, mock_run: MagicMock, mock_which: MagicMock) -> None:
+        """Tests behavior when no GPU is detected (nvidia-smi fails or missing)."""
+        # if nvidia-smi missing
+        mock_which.return_value = None
+        # On windows, it tries default path even if which returns None,
+        # so we also need subprocess to fail or return empty
+        mock_run.side_effect = FileNotFoundError("No file")
+
+        analyzer = GPUAnalyzer()
+        self.assertFalse(analyzer.is_gpu_available())
+        self.assertEqual(analyzer.gpu_count, 0)
+        self.assertEqual(analyzer.device_name, "CPU")
+        self.assertEqual(analyzer.vram, 0)
+        self.assertEqual(analyzer.get_device_name(0), "CPU")
+        self.assertEqual(analyzer.get_vram(0), 0)
+
+    @patch("canirun.gpu.shutil.which")
+    @patch("subprocess.run")
+    def test_multiple_gpus(self, mock_run: MagicMock, mock_which: MagicMock) -> None:
+        """Tests detection of multiple GPUs."""
+        mock_which.return_value = "nvidia-smi"
+        output = "0, GPU A, 8000\n" "1, GPU B, 12000"
+        mock_run.return_value.stdout = output
+        mock_run.return_value.returncode = 0
+
+        analyzer = GPUAnalyzer()
+
+        self.assertEqual(analyzer.gpu_count, 2)
+        self.assertEqual(analyzer.get_device_name(0), "GPU A")
+        self.assertEqual(analyzer.get_device_name(1), "GPU B")
+
+        # Test out of bounds
+        self.assertEqual(
+            analyzer.get_device_name(99), "GPU A"
+        )  # Should warn and return first
+        self.assertEqual(analyzer.get_vram(99), 8000 * 1024 * 1024)
+
+    @patch("subprocess.run")
+    def test_windows_default_path(self, mock_run: MagicMock) -> None:
+        """Tests fallback to default Windows path if shutil.which fails."""
+        with (
+            patch("sys.platform", "win32"),
+            patch("canirun.gpu.shutil.which", return_value=None),
+            patch("os.environ.get", return_value="C:"),
+        ):
+            mock_run.return_value.stdout = "0, GPU Default, 4000"
+            mock_run.return_value.returncode = 0
+
+            analyzer = GPUAnalyzer()
+
+            self.assertTrue(analyzer.is_gpu_available())
+            self.assertEqual(analyzer.device_name, "GPU Default")
+
+            # Verify it tried the hardcoded path
+            args, _ = mock_run.call_args
+            cmd = args[0]
+            self.assertTrue("Program Files" in cmd[0])
+            self.assertTrue("nvidia-smi.exe" in cmd[0])
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_logic.py

@@ -73,6 +73,78 @@ def test_fetch_model_data_auth_error(
             )
             self.assertTrue(found_tip, f"Tip not found in logs: {cm.output}")
 
+    @patch("canirun.logic.GPUAnalyzer")
+    @patch("psutil.virtual_memory")
+    def test_get_specs_with_gpu(self, mock_vm: Any, MockGPU: Any) -> None:
+        """Tests that _get_specs correctly prioritizes GPU over CPU/Mac."""
+        # Setup RAM
+        mock_vm.return_value.total = 32 * 1024**3
+
+        # Setup GPU
+        mock_gpu_instance = MockGPU.return_value
+        mock_gpu_instance.is_gpu_available.return_value = True
+        mock_gpu_instance.vram = 24 * 1024**3
+        mock_gpu_instance.device_name = "NVIDIA RTX 3090"
+
+        # Initialize analyzer
+        analyzer = ModelAnalyzer("test-model", verbose=False)
+
+        self.assertEqual(analyzer.specs["vram"], 24 * 1024**3)
+        self.assertEqual(analyzer.specs["name"], "NVIDIA RTX 3090")
+        self.assertFalse(analyzer.specs["is_mac"])
+
+    @patch("canirun.logic.GPUAnalyzer")
+    @patch("platform.machine")
+    @patch("platform.system")
+    @patch("psutil.virtual_memory")
+    def test_get_specs_mac_silicon(
+        self, mock_vm: Any, mock_system: Any, mock_machine: Any, MockGPU: Any
+    ) -> None:
+        """Tests that _get_specs correctly detects Apple Silicon."""
+        # Setup RAM
+        mock_vm.return_value.total = 16 * 1024**3
+
+        # Setup Mac Environment
+        mock_system.return_value = "Darwin"
+        mock_machine.return_value = "arm64"
+
+        # Setup GPU (No discrete GPU)
+        mock_gpu_instance = MockGPU.return_value
+        mock_gpu_instance.is_gpu_available.return_value = False
+
+        analyzer = ModelAnalyzer("test-model", verbose=False)
+
+        self.assertTrue(analyzer.specs["is_mac"])
+        self.assertEqual(analyzer.specs["name"], "Apple Silicon (Unified Memory)")
+        # Check VRAM calculation (75% of RAM)
+        expected_vram = 16 * 1024**3 * 0.75
+        self.assertEqual(analyzer.specs["vram"], expected_vram)
+
+    def test_calculate_fallback_params(self) -> None:
+        """Tests calculation fallback when params_billions is 0."""
+        # Data with 0 params but architecture details
+        data = {
+            "params_billions": 0,
+            "hidden_size": 1024,
+            "num_hidden_layers": 10,
+            "num_attention_heads": 8,
+            "num_key_value_heads": 8,
+            "vocab_size": 1000,
+        }
+
+        results = self.analyzer.calculate(data, ctx=1024)
+        self.assertTrue(len(results) > 0)
+
+        # Verify params were calculated
+        # block = 12 * 10 * 1024^2 = 120 * 1,048,576 = 125M
+        # embed = 1000 * 1024 = 1M
+        # Total = 126M params -> 0.126B
+        # FP16 size = 0.25 GB
+
+        res_fp16 = results[0]
+        self.assertEqual(res_fp16["quant"], "FP16")
+        self.assertGreater(res_fp16["total_ram"], 0)
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_utils.py

@@ -0,0 +1,31 @@
+import unittest
+
+from canirun.enum import COMPATIBILITY
+from canirun.human_readable import get_human_readable_size, get_human_readable_status
+
+
+class TestHumanReadable(unittest.TestCase):
+    """Test suite for human-readable helper functions."""
+
+    def test_get_human_readable_size(self) -> None:
+        """Tests bytes to human-readable string conversion."""
+        self.assertEqual(get_human_readable_size(0), "0 B")
+        self.assertEqual(get_human_readable_size(100), "100.00 B")
+        self.assertEqual(get_human_readable_size(1024), "1.00 KB")
+        self.assertEqual(get_human_readable_size(1024**2), "1.00 MB")
+        self.assertEqual(get_human_readable_size(1024**3), "1.00 GB")
+        self.assertEqual(get_human_readable_size(1.5 * 1024**3), "1.50 GB")
+        self.assertEqual(get_human_readable_size(1024**4), "1.00 TB")
+        self.assertEqual(get_human_readable_size(1024**5), "1.00 PB")
+
+    def test_get_human_readable_status(self) -> None:
+        """Tests compatibility enum to string conversion."""
+        self.assertEqual(get_human_readable_status(COMPATIBILITY.FULL), "✅ GPU")
+        self.assertEqual(
+            get_human_readable_status(COMPATIBILITY.PARTIAL), "⚠️ CPU/RAM only (Slow)"
+        )
+        self.assertEqual(get_human_readable_status(COMPATIBILITY.NONE), "❌ Impossible")
+
+
+if __name__ == "__main__":
+    unittest.main()