Refactor TabPFN tests to use patch decorators and add integration test

- Replace global `sys.modules` mocking of `tabpfn` with `@patch` decorators in `TestTabPFNClassifierClass` for better test isolation.
- Update test methods (`test_fit_v2_5_default`, `test_fit_v2_5_synthetic`, `test_fit_v2`, etc.) to accept and use the patched mock object directly.
- Add `test_real_execution_if_available` to verify integration with the actual `tabpfn` library if installed, gracefully handling model download/gating errors.
- Simplify `setUp` method by removing manual mock resets.
- Improve conditional import logic for `TabPFNClassifierClass`.

Author: SamoraHunter

tests/test_tabpfn_classifier_class.py

@@ -3,6 +3,9 @@
 import os
 import pandas as pd
 from unittest.mock import MagicMock, patch
+from ml_grid.model_classes.tabpfn_classifier_class import (
+    TabPFNClassifierClass,
+)
 
 # Add project root to path so we can import the class under test
 # Assumes this test file is in /tests/ and ml_grid is in the parent dir
@@ -22,37 +25,18 @@
 sys.modules["ml_grid.util.param_space"] = MagicMock()
 sys.modules["ml_grid.util.global_params"] = MagicMock()
 
-# Mock tabpfn to avoid downloading weights during tests
-mock_tabpfn_module = MagicMock()
-# The mock needs to have the attributes that are accessed in the class
-mock_tabpfn_classifier_instance = MagicMock()
-mock_tabpfn_module.TabPFNClassifier.return_value = mock_tabpfn_classifier_instance
-mock_tabpfn_module.TabPFNClassifier.create_default_for_version.return_value = (
-    mock_tabpfn_classifier_instance
-)
-sys.modules["tabpfn"] = mock_tabpfn_module
-sys.modules["tabpfn.constants"] = MagicMock()
-
-from ml_grid.model_classes.tabpfn_classifier_class import (
-    TabPFNClassifierClass,
-)  # noqa: E402
+# Conditionally mock tabpfn. If it's installed, we might want to use it for integration tests.
+# If not installed, we mock it to allow importing the wrapper class.
+try:
+    import tabpfn
+except ImportError:
+    mock_tabpfn_module = MagicMock()
+    sys.modules["tabpfn"] = mock_tabpfn_module
+    sys.modules["tabpfn.constants"] = MagicMock()
 
 
 class TestTabPFNClassifierClass(unittest.TestCase):
     def setUp(self):
-        # Reset the mock for TabPFNClassifier before each test
-        mock_tabpfn_module.TabPFNClassifier.reset_mock()
-        mock_tabpfn_module.TabPFNClassifier.create_default_for_version.reset_mock()
-        # also reset the instance mock
-        mock_tabpfn_classifier_instance.reset_mock()
-        # re-assign the return value in case it was modified
-        mock_tabpfn_module.TabPFNClassifier.return_value = (
-            mock_tabpfn_classifier_instance
-        )
-        mock_tabpfn_module.TabPFNClassifier.create_default_for_version.return_value = (
-            mock_tabpfn_classifier_instance
-        )
-
         # Patch global parameters to control bayessearch flag
         self.global_params_patch = patch("ml_grid.util.global_params.global_parameters")
         self.mock_global_params = self.global_params_patch.start()
@@ -98,8 +82,12 @@ def test_parameter_space_compatibility(self):
         for param in removed_params:
             self.assertNotIn(param, model.parameter_space)
 
-    def test_fit_v2_5_default(self):
+    @patch("ml_grid.model_classes.tabpfn_classifier_class.TabPFNClassifier")
+    def test_fit_v2_5_default(self, mock_tabpfn_cls):
         """Test fitting of the default v2.5 model."""
+        # Setup mock return value
+        mock_estimator_instance = mock_tabpfn_cls.return_value
+
         # Instantiate the wrapper with hyperparameters
         model_wrapper = TabPFNClassifierClass(
             model_version="v2.5_default", n_estimators=4, device="cpu", random_state=42
@@ -109,10 +97,10 @@ def test_fit_v2_5_default(self):
         model_wrapper.fit(self.X_dummy, self.y_dummy)
 
         # Verify TabPFNClassifier constructor was called directly
-        mock_tabpfn_module.TabPFNClassifier.assert_called_once()
+        mock_tabpfn_cls.assert_called_once()
 
         # Verify arguments passed to the constructor
-        call_kwargs = mock_tabpfn_module.TabPFNClassifier.call_args.kwargs
+        call_kwargs = mock_tabpfn_cls.call_args.kwargs
         self.assertEqual(call_kwargs["n_estimators"], 4)
         self.assertEqual(call_kwargs["device"], "cpu")
         self.assertEqual(call_kwargs["random_state"], 42)
@@ -122,48 +110,56 @@ def test_fit_v2_5_default(self):
         self.assertNotIn("subsample_samples", call_kwargs)
 
         # Verify the underlying estimator's fit method was called
-        mock_estimator_instance = mock_tabpfn_module.TabPFNClassifier.return_value
         mock_estimator_instance.fit.assert_called_once()
 
-    def test_fit_v2_5_synthetic(self):
+    @patch("ml_grid.model_classes.tabpfn_classifier_class.TabPFNClassifier")
+    def test_fit_v2_5_synthetic(self, mock_tabpfn_cls):
         """Test fitting of the synthetic v2.5 model (checks model_path logic)."""
+        mock_estimator_instance = mock_tabpfn_cls.return_value
+
         model_wrapper = TabPFNClassifierClass(
             model_version="v2.5_synthetic", n_estimators=2
         )
 
         model_wrapper.fit(self.X_dummy, self.y_dummy)
 
         # Verify constructor was called
-        mock_tabpfn_module.TabPFNClassifier.assert_called_once()
+        mock_tabpfn_cls.assert_called_once()
 
         # Verify model_path was injected
-        call_kwargs = mock_tabpfn_module.TabPFNClassifier.call_args.kwargs
+        call_kwargs = mock_tabpfn_cls.call_args.kwargs
         self.assertEqual(
             call_kwargs.get("model_path"), "tabpfn-v2.5-classifier-v2.5_default-2.ckpt"
         )
 
         # Verify fit was called
-        mock_estimator_instance = mock_tabpfn_module.TabPFNClassifier.return_value
         mock_estimator_instance.fit.assert_called_once()
 
-    def test_fit_v2(self):
+    @patch("ml_grid.model_classes.tabpfn_classifier_class.TabPFNClassifier")
+    def test_fit_v2(self, mock_tabpfn_cls):
         """Test fitting of the legacy v2 model."""
+        # Setup mock for create_default_for_version
+        mock_estimator_instance = MagicMock()
+        mock_tabpfn_cls.create_default_for_version.return_value = (
+            mock_estimator_instance
+        )
+
         model_wrapper = TabPFNClassifierClass(model_version="v2", n_estimators=1)
 
         model_wrapper.fit(self.X_dummy, self.y_dummy)
 
         # Verify it called create_default_for_version instead of standard constructor
-        mock_tabpfn_module.TabPFNClassifier.create_default_for_version.assert_called_once()
-        mock_tabpfn_module.TabPFNClassifier.assert_not_called()  # Ensure standard constructor was NOT called
+        mock_tabpfn_cls.create_default_for_version.assert_called_once()
+        mock_tabpfn_cls.assert_not_called()  # Ensure standard constructor was NOT called
 
         # Verify fit was called
-        mock_estimator_instance = (
-            mock_tabpfn_module.TabPFNClassifier.create_default_for_version.return_value
-        )
         mock_estimator_instance.fit.assert_called_once()
 
-    def test_predict_and_predict_proba_delegation(self):
+    @patch("ml_grid.model_classes.tabpfn_classifier_class.TabPFNClassifier")
+    def test_predict_and_predict_proba_delegation(self, mock_tabpfn_cls):
         """Test that predict and predict_proba delegate to the internal estimator."""
+        mock_estimator_instance = mock_tabpfn_cls.return_value
+
         model_wrapper = TabPFNClassifierClass()
 
         # Fit the model to create the internal _estimator
@@ -180,8 +176,11 @@ def test_predict_and_predict_proba_delegation(self):
         model_wrapper.predict_proba(self.X_dummy)
         internal_estimator_mock.predict_proba.assert_called_once_with(self.X_dummy)
 
-    def test_fit_with_subsampling(self):
+    @patch("ml_grid.model_classes.tabpfn_classifier_class.TabPFNClassifier")
+    def test_fit_with_subsampling(self, mock_tabpfn_cls):
         """Test that subsampling is applied when configured."""
+        mock_estimator_instance = mock_tabpfn_cls.return_value
+
         # Create larger dummy data
         X_large = pd.DataFrame({"col1": range(100), "col2": range(100)})
         y_large = pd.Series([0, 1] * 50)
@@ -194,17 +193,53 @@ def test_fit_with_subsampling(self):
         model_wrapper.fit(X_large, y_large)
 
         # Verify constructor called without subsample_samples
-        call_kwargs = mock_tabpfn_module.TabPFNClassifier.call_args.kwargs
+        call_kwargs = mock_tabpfn_cls.call_args.kwargs
         self.assertNotIn("subsample_samples", call_kwargs)
 
         # Verify fit was called with subsampled data
-        mock_estimator_instance = mock_tabpfn_module.TabPFNClassifier.return_value
         args, _ = mock_estimator_instance.fit.call_args
         X_passed, y_passed = args
 
         self.assertEqual(len(X_passed), subsample_size)
         self.assertEqual(len(y_passed), subsample_size)
 
+    def test_real_execution_if_available(self):
+        """
+        Integration test: Attempts to run with the real TabPFN library if available.
+        If the model weights are missing (gated), it catches the RuntimeError and passes.
+        """
+        try:
+            # Try to instantiate and fit a small model
+            # We use n_estimators=1 for speed
+            model_wrapper = TabPFNClassifierClass(n_estimators=1, device="cpu")
+            model_wrapper.fit(self.X_dummy, self.y_dummy)
+
+            # If fit succeeds, try predict
+            preds = model_wrapper.predict(self.X_dummy)
+            self.assertEqual(len(preds), len(self.X_dummy))
+
+        except RuntimeError as e:
+            # Check for the specific download error
+            error_msg = str(e).lower()
+            if (
+                "download" in error_msg
+                or "gated" in error_msg
+                or "modelversion" in error_msg
+            ):
+                print(
+                    f"Skipping real execution test (Model download required/gated): {e}"
+                )
+                return
+            # If it's another RuntimeError, re-raise it
+            raise e
+        except ImportError:
+            print("Skipping real execution test (TabPFN not installed)")
+            return
+        except Exception as e:
+            # Catch-all for other environment issues (e.g. network)
+            print(f"Skipping real execution test due to unexpected error: {e}")
+            return
+
 
 if __name__ == "__main__":
     unittest.main()