test: eval env builder

tests/test_evaluation.py

@@ -14,6 +14,7 @@ from brittle_star_project.environment.env_config import (
     ArenaConfig,
     EnvConfig,
     ObservationBoundsConfig,
+    MorphMode,
 )
 from brittle_star_project.environment.env_types import Task
 
@@ -112,3 +113,106 @@ def test_load_metadata_with_override(tmp_path: Path):
     non_existent = tmp_path / "missing.yaml"
     with pytest.raises(FileNotFoundError, match="Could not find metadata YAML at"):
         load_metadata(model_path, metadata_override_path=non_existent)
+
+
+@pytest.fixture
+def mock_training_config():
+    return TrainingConfig(
+        morphology=MorphologyConfig(
+            segments_per_arm=[1, 1, 1, 1, 1], morph_mode=MorphMode.CENTRALIZED
+        ),
+        arena=ArenaConfig(),
+        environment=EnvConfig(),
+        obs_bounds=ObservationBoundsConfig(),
+    )
+
+
+@pytest.fixture
+def mock_metadata():
+    return {"architecture": {"message_passing_steps": 2}}
+
+
+def test_build_eval_env_training_morphology(tmp_path, mock_training_config, mock_metadata):
+    from brittle_star_project.evaluation.eval_env_builder import build_eval_env
+    from unittest.mock import patch
+
+    model_path = tmp_path / "model.flax"
+    patch_target = "brittle_star_project.evaluation.eval_env_builder.PolicyAgent.from_checkpoint"
+    with patch(patch_target) as mock_agent:
+        mock_agent.return_value = "mock_policy"
+        bundle = build_eval_env(
+            model_path=model_path,
+            training=mock_training_config,
+            metadata=mock_metadata,
+            morphology_override_path=None,
+        )
+        assert bundle.segments_per_arm == [1, 1, 1, 1, 1]
+        assert bundle.num_active_arms == 5
+        assert bundle.architecture == "CENTRALIZED"
+        assert bundle.policy == "mock_policy"
+
+
+def test_build_eval_env_override_morphology(tmp_path, mock_training_config, mock_metadata):
+    from brittle_star_project.evaluation.eval_env_builder import build_eval_env
+    from unittest.mock import patch
+
+    model_path = tmp_path / "model.flax"
+    override_path = tmp_path / "override.yaml"
+    override_path.write_text(yaml.dump({"segments_per_arm": [1, 0, 1, 0, 1]}))
+
+    with patch("brittle_star_project.evaluation.eval_env_builder.PolicyAgent.from_checkpoint"):
+        bundle = build_eval_env(
+            model_path=model_path,
+            training=mock_training_config,
+            metadata=mock_metadata,
+            morphology_override_path=override_path,
+        )
+        assert bundle.segments_per_arm == [1, 0, 1, 0, 1]
+        assert bundle.num_active_arms == 3
+        # Should be smaller than 5*N
+        assert sum(bundle.action_mask) < len(bundle.action_mask)
+
+
+def test_build_eval_env_action_mask_shape(tmp_path, mock_training_config, mock_metadata):
+    from brittle_star_project.evaluation.eval_env_builder import build_eval_env
+    from unittest.mock import patch
+
+    model_path = tmp_path / "model.flax"
+    override_path = tmp_path / "override.yaml"
+    override_path.write_text(yaml.dump({"segments_per_arm": [1, 0, 1, 0, 0]}))
+
+    with patch("brittle_star_project.evaluation.eval_env_builder.PolicyAgent.from_checkpoint"):
+        bundle = build_eval_env(
+            model_path=model_path,
+            training=mock_training_config,
+            metadata=mock_metadata,
+            morphology_override_path=override_path,
+        )
+        # For each segment with P-control, there's 2 actions (pitch and yaw).
+        # Total segments = 5 -> 10 actions for training.
+        assert len(bundle.action_mask) == 10
+        # Active segments = 2 -> 4 actions active.
+        assert sum(bundle.action_mask) == 4
+
+
+def test_build_eval_env_morph_mode_inherited(tmp_path, mock_training_config, mock_metadata):
+    from brittle_star_project.evaluation.eval_env_builder import build_eval_env
+    from brittle_star_project.environment.env_config import MorphMode
+    from unittest.mock import patch
+
+    model_path = tmp_path / "model.flax"
+    override_path = tmp_path / "override.yaml"
+    # No morph_mode in the override YAML
+    override_path.write_text(yaml.dump({"segments_per_arm": [1, 0, 1, 0, 1]}))
+
+    # Change training config to be RING
+    mock_training_config.morphology.morph_mode = MorphMode.RING
+
+    with patch("brittle_star_project.evaluation.eval_env_builder.PolicyAgent.from_checkpoint"):
+        bundle = build_eval_env(
+            model_path=model_path,
+            training=mock_training_config,
+            metadata=mock_metadata,
+            morphology_override_path=override_path,
+        )
+        assert bundle.architecture == "RING"