No more test mocking (#58)

Author: msaroufim

test/test_backends.py

@@ -1,5 +1,3 @@
-from unittest.mock import Mock, patch
-
 import pytest
 import torch
 from BackendBench.backends import (
@@ -38,128 +36,99 @@ def test_aten_backend_contains_op(self):
 
         assert torch.ops.aten.relu.default in backend
         assert torch.ops.aten.add.Tensor in backend
-
-        fake_op = Mock()
-        fake_op.__module__ = "fake_module"
-        assert fake_op in backend  # AtenBackend contains everything
+        assert torch.ops.aten.mul.Tensor in backend
 
     def test_aten_backend_getitem(self):
         backend = AtenBackend()
 
         relu_op = torch.ops.aten.relu.default
         assert backend[relu_op] == relu_op
 
-        fake_op = Mock()
-        fake_op.__module__ = "fake_module"
-        assert backend[fake_op] == fake_op  # AtenBackend returns the op itself
+        add_op = torch.ops.aten.add.Tensor
+        assert backend[add_op] == add_op
 
 
 class TestFlagGemsBackend:
     @pytest.mark.skipif(not HAS_FLAG_GEMS, reason="flag_gems not available")
-    @patch("BackendBench.backends.flag_gems")
-    def test_flag_gems_backend_initialization(self, mock_flag_gems):
+    def test_flag_gems_backend_initialization(self):
         backend = FlagGemsBackend()
         assert backend.name == "flaggems"
         assert isinstance(backend.ops, dict)
 
     @pytest.mark.skipif(not HAS_FLAG_GEMS, reason="flag_gems not available")
-    @patch("BackendBench.backends.flag_gems")
-    def test_flag_gems_backend_contains_op(self, mock_flag_gems):
-        mock_flag_gems.abs = Mock()
-
+    def test_flag_gems_backend_contains_op(self):
         backend = FlagGemsBackend()
 
-        assert torch.ops.aten.abs.default in backend
+        # Test with actual ops that flag_gems supports
+        if hasattr(torch.ops.aten, "abs"):
+            if torch.ops.aten.abs.default in backend:
+                assert torch.ops.aten.abs.default in backend
 
-        fake_op = Mock()
-        fake_op.__str__ = Mock(return_value="fake_op")
-        assert fake_op not in backend
+        # Test with an op that might not be in flag_gems
+        unsupported_op = (
+            torch.ops.aten.special_log_ndtr.default
+            if hasattr(torch.ops.aten, "special_log_ndtr")
+            else None
+        )
+        if unsupported_op:
+            assert unsupported_op not in backend
 
     @pytest.mark.skipif(not HAS_FLAG_GEMS, reason="flag_gems not available")
-    @patch("BackendBench.backends.flag_gems")
-    def test_flag_gems_backend_getitem(self, mock_flag_gems):
-        mock_abs_impl = Mock()
-        mock_flag_gems.ops.abs = mock_abs_impl
-
+    def test_flag_gems_backend_getitem(self):
         backend = FlagGemsBackend()
 
-        assert backend[torch.ops.aten.abs.default] == mock_abs_impl
+        # Test with an op that should exist
+        if hasattr(torch.ops.aten, "abs") and torch.ops.aten.abs.default in backend:
+            impl = backend[torch.ops.aten.abs.default]
+            assert impl is not None
 
-        fake_op = Mock()
-        fake_op.__str__ = Mock(return_value="fake_op")
-        with pytest.raises(KeyError):
-            _ = backend[fake_op]
+        # Test with an op that doesn't exist in flag_gems
+        unsupported_op = (
+            torch.ops.aten.special_log_ndtr.default
+            if hasattr(torch.ops.aten, "special_log_ndtr")
+            else None
+        )
+        if unsupported_op and unsupported_op not in backend:
+            with pytest.raises(KeyError):
+                _ = backend[unsupported_op]
 
 
 class TestLLMBackend:
     def test_llm_backend_initialization(self):
-        with (
-            patch("os.makedirs"),
-            patch("builtins.open"),
-            patch("datetime.datetime") as mock_datetime,
-        ):
-            mock_datetime.now.return_value.strftime.return_value = "20250721_204542"
-            backend = LLMBackend()
-            assert backend.name == "llm"
-            assert "generated_kernels/run_" in backend.kernels_dir
-            assert isinstance(backend.compiled_kernels, dict)
-
-    @pytest.mark.skip(
-        reason="Complex file I/O mocking needed - test requires full file system interaction"
-    )
-    def test_llm_backend_add_kernel(self):
-        with (
-            patch("os.makedirs"),
-            patch("builtins.open"),
-            patch("datetime.datetime") as mock_datetime,
-        ):
-            mock_datetime.now.return_value.strftime.return_value = "20250721_204542"
-            backend = LLMBackend()
-
-            mock_op = Mock()
-            mock_op.__name__ = "test_op"
-
-            kernel_code = """
-def test_kernel(x):
-    return x + 1
-"""
-
-            with patch("builtins.open", create=True) as mock_open:
-                backend.add_kernel(mock_op, kernel_code, "test_op")
-
-            mock_open.assert_called()
-
-            assert mock_op in backend
-
-    @pytest.mark.skip(
-        reason="Complex file I/O mocking needed - test requires full file system interaction"
-    )
-    def test_llm_backend_test_kernel_correctness(self):
-        with (
-            patch("os.makedirs"),
-            patch("builtins.open"),
-            patch("datetime.datetime") as mock_datetime,
-        ):
-            mock_datetime.now.return_value.strftime.return_value = "20250721_204542"
-            backend = LLMBackend()
+        backend = LLMBackend()
+        assert backend.name == "llm"
+        assert "generated_kernels/run_" in backend.kernels_dir
+        assert isinstance(backend.compiled_kernels, dict)
 
-            mock_op = Mock(return_value=torch.tensor([2.0]))
-
-            kernel_code = """
-def generated_kernel(x):
-    return x + 1
+    @pytest.mark.skip(reason="Requires Triton for kernel compilation")
+    def test_llm_backend_add_kernel(self):
+        backend = LLMBackend()
+
+        # Use a real torch op for testing
+        test_op = torch.ops.aten.relu.default
+
+        kernel_code = """
+@triton.jit
+def relu_kernel(x_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
+    pid = tl.program_id(0)
+    block_start = pid * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+    x = tl.load(x_ptr + offsets, mask=mask)
+    output = tl.maximum(x, 0)
+    tl.store(output_ptr + offsets, output, mask=mask)
+
+def generated_relu(x):
+    output = torch.empty_like(x)
+    n_elements = output.numel()
+    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
+    relu_kernel[grid](x, output, n_elements, BLOCK_SIZE=1024)
+    return output
 """
 
-            mock_test = Mock()
-            mock_test.args = [torch.tensor([1.0])]
-            mock_test.kwargs = {}
-
-            with patch("builtins.open", create=True):
-                is_correct, feedback = backend.test_kernel_correctness(
-                    mock_op, kernel_code, [mock_test], attempt=1
-                )
+        backend.add_kernel(test_op, kernel_code, "relu")
 
-            assert is_correct is True
+        assert test_op in backend
 
 
 class TestKernelAgentBackend:
@@ -180,42 +149,20 @@ def test_kernel_agent_backend_set_config(self):
         assert backend.num_workers == 8
         assert backend.max_rounds == 20
 
-    @pytest.mark.skipif(not HAS_KERNEL_AGENT, reason="KernelAgent not available")
-    def test_kernel_agent_backend_generate_kernel(self):
-        with (
-            patch("triton_kernel_agent.TritonKernelAgent") as mock_kernel_agent_class,
-        ):
-            backend = KernelAgentBackend()
-
-            mock_agent = Mock()
-            mock_kernel_agent_class.return_value = mock_agent
-
-            mock_agent.generate_kernel.return_value = {
-                "success": True,
-                "kernel_code": "def kernel(): pass",
-                "rounds": 1,
-                "session_dir": "test_session_dir",
-                "worker_id": 0,
-            }
-
-            mock_op = Mock()
-            mock_op.__str__ = Mock(return_value="test_op")
-            with patch("builtins.open", create=True):
-                kernel_code, success = backend.generate_kernel_with_agent(mock_op, "test_op")
-            assert success is True
-            assert kernel_code == "def kernel(): pass"
-            mock_kernel_agent_class.assert_called_once()
-
 
 class TestBackendIntegration:
-    @pytest.mark.skipif(not HAS_FLAG_GEMS, reason="flag_gems not available")
     def test_backend_polymorphism(self):
         backends = []
         backends.append(AtenBackend())
-        with patch("BackendBench.backends.flag_gems"):
+
+        if HAS_FLAG_GEMS:
             backends.append(FlagGemsBackend())
+
         backends.append(LLMBackend())
-        backends.append(KernelAgentBackend())
+
+        if HAS_KERNEL_AGENT:
+            backends.append(KernelAgentBackend())
+
         for backend in backends:
             assert hasattr(backend, "name")
             assert hasattr(backend, "__contains__")

test/test_directory_backend.py

@@ -14,9 +14,10 @@
 
 @pytest.fixture(scope="module")
 def backend():
-    # Ensure generated_kernels directory exists for CI
-    if not os.path.exists("generated_kernels"):
-        # Import and run the existing script
+    expected_dirs = ["relu", "add", "mul", "abs", "sum"]
+    missing_dirs = [d for d in expected_dirs if not os.path.isdir(f"generated_kernels/{d}")]
+
+    if missing_dirs:
         import subprocess
 
         subprocess.run(