Update tests

Author: toby-coleman

tests/unit/test_tuner.py

@@ -1,9 +1,12 @@
 """Provides unit tests for the Tuner class."""
 
+from tempfile import TemporaryDirectory
+import typing as _t
 from unittest.mock import MagicMock, patch
 
 import msgspec
 import pytest
+import ray.tune
 
 from plugboard.schemas import ConfigSpec, ObjectiveSpec
 from plugboard.schemas.tune import (
@@ -23,6 +26,13 @@ def config() -> dict:
         return msgspec.yaml.decode(f.read())
 
 
+@pytest.fixture
+def temp_dir() -> _t.Iterator[str]:
+    """Creates a temporary directory."""
+    with TemporaryDirectory() as tmpdir:
+        yield tmpdir
+
+
 @patch("ray.tune.Tuner")
 def test_tuner(mock_tuner_cls: MagicMock, config: dict) -> None:
     """Test the Tuner class."""
@@ -91,64 +101,7 @@ def test_tuner(mock_tuner_cls: MagicMock, config: dict) -> None:
     mock_tuner.fit.assert_called_once()
 
 
-@patch("ray.tune.Tuner")
-def test_tuner_with_optuna_storage(mock_tuner_cls: MagicMock, config: dict) -> None:
-    """Test the Tuner class with Optuna storage URI."""
-    mock_tuner = MagicMock()
-    mock_tuner_cls.return_value = mock_tuner
-
-    spec = ConfigSpec.model_validate(config)
-    process_spec = spec.plugboard.process
-
-    # Test with storage URI
-    optuna_spec = OptunaSpec(
-        type="ray.tune.search.optuna.OptunaSearch",
-        study_name="test-study",
-        storage="sqlite:///test.db",
-    )
-
-    tuner = Tuner(
-        objective=ObjectiveSpec(
-            object_type="component",
-            object_name="c",
-            field_type="field",
-            field_name="in_1",
-        ),
-        parameters=[
-            FloatParameterSpec(
-                object_type="component",
-                object_name="a",
-                field_type="arg",
-                field_name="y",
-                lower=0.1,
-                upper=0.5,
-            ),
-        ],
-        num_samples=6,
-        mode="max",
-        max_concurrent=2,
-        algorithm=optuna_spec,
-    )
-    tuner.run(spec=process_spec)
-
-    # Must call the Tuner class with objective
-    assert callable(mock_tuner_cls.call_args.args[0])
-    # Must call the Tuner class with parameter space
-    kwargs = mock_tuner_cls.call_args.kwargs
-    param_space = kwargs["param_space"]
-    assert param_space["a.y"].__class__.__name__ == "Float"
-    assert param_space["a.y"].lower == 0.1
-    assert param_space["a.y"].upper == 0.5
-    # Must call the Tuner class with configuration and correct algorithm
-    tune_config = kwargs["tune_config"]
-    assert tune_config.num_samples == 6
-    # Check searcher attribute as this contains the underlying algorithm with storage converted
-    assert tune_config.search_alg.searcher.__class__.__name__ == "OptunaSearch"
-    # Must call fit method on the Tuner object
-    mock_tuner.fit.assert_called_once()
-
-
-def test_optuna_storage_uri_conversion() -> None:
+def test_optuna_storage_uri_conversion(temp_dir: str) -> None:
     """Test that storage URI gets converted to Optuna storage object."""
     # Create a tuner with minimal configuration
     tuner = Tuner(
@@ -167,19 +120,11 @@ def test_optuna_storage_uri_conversion() -> None:
         ],
         num_samples=1,
         mode="max",
+        algorithm=OptunaSpec(
+            type="ray.tune.search.optuna.OptunaSearch",
+            study_name="test-study",
+            storage=f"sqlite:///{temp_dir}/test_conversion.db",
+        ),
     )
-
-    # Test the _build_algorithm method with storage URI
-    optuna_spec = OptunaSpec(
-        type="ray.tune.search.optuna.OptunaSearch",
-        study_name="test-study",
-        storage="sqlite:///test_conversion.db",
-    )
-
-    # This should work without raising AssertionError
-    algorithm = tuner._build_algorithm(optuna_spec)
-
-    # Verify the algorithm was created successfully
-    import ray.tune.search.optuna
-
-    assert isinstance(algorithm, ray.tune.search.optuna.OptunaSearch)
+    algo = tuner._config.search_alg
+    assert isinstance(algo, ray.tune.search.optuna.OptunaSearch)