File indexing completed on 2026-08-12 08:24:58
0001 """Tests for the canonical optimizer config registry."""
0002
0003 from __future__ import annotations
0004
0005 import sys
0006 from importlib import import_module
0007
0008 import pytest
0009 from pydantic import BaseModel
0010
0011 from aid2e.utilities.configurations.optimization_registry import (
0012 get,
0013 list_registered,
0014 register,
0015 )
0016 from aid2e.utilities.configurations.optimizer_config import OptimizerConfiguration
0017
0018
0019 @pytest.fixture
0020 def clean_optimizer_registry(monkeypatch):
0021 """Reset registry state so lazy loading behavior is exercised explicitly."""
0022 registry_module = import_module("aid2e.utilities.configurations.optimization_registry")
0023 monkeypatch.setattr(registry_module, "_algorithm_configs", {})
0024
0025 for module_name in [
0026 "aid2e.optimizers",
0027 "aid2e.optimizers.ax",
0028 "aid2e.optimizers.ax.config",
0029 "aid2e.optimizers.ax.optimizer",
0030 "aid2e.optimizers.pymoo",
0031 "aid2e.optimizers.pymoo.config",
0032 "aid2e.optimizers.pymoo.optimizer",
0033 ]:
0034 sys.modules.pop(module_name, None)
0035
0036 return registry_module
0037
0038
0039 def test_get_lazily_loads_ax_config(clean_optimizer_registry) -> None:
0040 model = get("ax")
0041
0042 assert model is not None
0043 assert model.__name__ == "AxOptimizerConfig"
0044
0045
0046 def test_get_lazily_loads_pymoo_config(clean_optimizer_registry) -> None:
0047 model = get("pymoo")
0048
0049 assert model is not None
0050 assert model.__name__ == "PyMOOOptimizerConfig"
0051
0052
0053 def test_parse_algorithm_params_returns_ax_model(clean_optimizer_registry) -> None:
0054 config = OptimizerConfiguration(
0055 name="ax",
0056 type="Bayesian",
0057 parameters={
0058 "generator": "BOTORCH_MODULAR",
0059 "n_iterations": 5,
0060 "batch_size": 2,
0061 },
0062 )
0063
0064 parsed = config.parse_algorithm_params()
0065
0066 assert parsed is not None
0067 assert parsed.__class__.__name__ == "AxOptimizerConfig"
0068 assert parsed.n_iterations == 5
0069 assert parsed.batch_size == 2
0070
0071
0072 def test_parse_algorithm_params_returns_pymoo_model(clean_optimizer_registry) -> None:
0073 config = OptimizerConfiguration(
0074 name="pymoo",
0075 type="evolutionary",
0076 parameters={
0077 "pop_size": 16,
0078 "n_offsprings": 8,
0079 "n_iterations": 5,
0080 },
0081 )
0082
0083 parsed = config.parse_algorithm_params()
0084
0085 assert parsed is not None
0086 assert parsed.__class__.__name__ == "PyMOOOptimizerConfig"
0087 assert parsed.pop_size == 16
0088 assert parsed.n_offsprings == 8
0089
0090
0091 def test_list_registered_includes_builtin_backends(clean_optimizer_registry) -> None:
0092 registered = list_registered()
0093
0094 assert "ax" in registered
0095 assert "pymoo" in registered
0096 assert registered["ax"].__name__ == "AxOptimizerConfig"
0097 assert registered["pymoo"].__name__ == "PyMOOOptimizerConfig"
0098
0099
0100 def test_register_rejects_duplicate_names(clean_optimizer_registry) -> None:
0101 class FirstConfig(BaseModel):
0102 pass
0103
0104 class SecondConfig(BaseModel):
0105 pass
0106
0107 register("custom", FirstConfig)
0108
0109 with pytest.raises(ValueError, match="already registered"):
0110 register("custom", SecondConfig)
