tests/test_environment.py

"""
test_environment.py — Tests for AdaptiveWorld environment.

Tests:
  1. Scenario registry loads correctly (all 12 scenarios present)
  2. DriftInjector — initial vs mutated world
  3. AdaptiveGrader — grade_task and grade_belief
  4. DriftDifficultyController — escalation logic
  5. Environment reset/step (manual reward signal verification)
  6. Drift detection after drift step
  7. Cross-episode belief persistence
  8. Expert: transient error simulation
  9. Grader fallback: infer_belief_from_actions
"""
import pytest
import sys
import os

# Add project root to path for imports
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from scenarios.registry import SCENARIO_REGISTRY, ALL_SCENARIOS
from server.drift_injector import DriftInjector, DRIFT_CONFIGS
from server.difficulty_controller import DriftDifficultyController
from graders.grader import AdaptiveGrader
from models import AdaptiveAction, AdaptiveObservation, AdaptiveState


# ── 1. Scenario registry ───────────────────────────────────────────────────────

class TestScenarioRegistry:

    def test_all_difficulty_levels_present(self):
        for level in ["easy", "medium", "hard", "expert"]:
            assert level in SCENARIO_REGISTRY, f"Missing difficulty: {level}"

    def test_each_level_has_three_scenarios(self):
        for level, scenarios in SCENARIO_REGISTRY.items():
            assert len(scenarios) == 3, f"{level} should have 3 scenarios, got {len(scenarios)}"

    def test_total_scenario_count(self):
        assert len(ALL_SCENARIOS) == 12

    def test_all_scenarios_have_required_fields(self):
        required = ["id", "description", "domain", "drift_trigger_step",
                    "drift_type", "task_goal", "max_steps"]
        for s in ALL_SCENARIOS:
            for field in required:
                assert field in s, f"Scenario {s.get('id')} missing field: {field}"

    def test_all_scenario_ids_in_drift_configs(self):
        """Every scenario must have a matching DriftInjector config."""
        for s in ALL_SCENARIOS:
            sid = s["id"]
            assert sid in DRIFT_CONFIGS, (
                f"Scenario '{sid}' has no entry in drift_injector.DRIFT_CONFIGS"
            )

    def test_no_drift_occurred_in_scenarios(self):
        """Drift provenance must be hidden — no drift_occurred / drift_type in scenario dict
        that would leak to the agent observation."""
        for s in ALL_SCENARIOS:
            # The key exists internally for grading but must NOT be a field name that
            # gets put into AdaptiveObservation (which no longer has drift_occurred)
            keys = list(s.keys())
            assert "drift_occurred" not in keys, (
                f"Scenario {s['id']} has drift_occurred — this should not be in observation"
            )


# ── 2. DriftInjector ──────────────────────────────────────────────────────────

class TestDriftInjector:

    def test_initial_world_loaded(self):
        inj = DriftInjector("easy_field_rename")
        world = inj.get_world()
        assert world["order_field"] == "qty"
        assert world.get("required_extra") is None

    def test_inject_changes_world(self):
        inj = DriftInjector("easy_field_rename")
        inj.inject()
        world = inj.get_world()
        assert world["order_field"] == "quantity"
        assert world["required_extra"] == "customer_id"

    def test_drifted_flag_set(self):
        inj = DriftInjector("easy_endpoint_version")
        assert not inj.drifted
        inj.inject()
        assert inj.drifted

    def test_world_truth_is_mutated_state(self):
        inj = DriftInjector("hard_status_meaning")
        truth = inj.get_world_truth()
        assert truth["order_status"] == "approved"  # mutated value

    def test_expert_transient_error_step(self):
        inj = DriftInjector("expert_transient_vs_real")
        assert inj.get_transient_error_step() == 2

    def test_expert_secondary_drift_step(self):
        inj = DriftInjector("expert_cross_service")
        assert inj.get_secondary_drift_step() == 6

    def test_unknown_scenario_raises(self):
        with pytest.raises(ValueError, match="Unknown scenario"):
            DriftInjector("nonexistent_scenario")


# ── 3. AdaptiveGrader ─────────────────────────────────────────────────────────

class TestAdaptiveGrader:
    grader = AdaptiveGrader()

    # grade_task tests
    def test_task_reward_failed_no_detection(self):
        r = self.grader.grade_task(
            task_completed=False, steps_taken=5, max_steps=8, drift_detected=False
        )
        assert r == 0.001

    def test_task_reward_failed_with_detection(self):
        r = self.grader.grade_task(
            task_completed=False, steps_taken=5, max_steps=8, drift_detected=True
        )
        assert r == 0.150

    def test_task_reward_completed_basic(self):
        r = self.grader.grade_task(
            task_completed=True, steps_taken=4, max_steps=8, drift_detected=False
        )
        assert 0.5 <= r <= 0.999

    def test_task_reward_completed_proactive_bonus(self):
        r_no_bonus = self.grader.grade_task(True, 4, 8, False)
        r_bonus = self.grader.grade_task(True, 4, 8, True)
        assert r_bonus > r_no_bonus

    def test_task_reward_capped_at_1(self):
        r = self.grader.grade_task(True, 1, 8, True)
        assert r <= 0.999

    # grade_belief tests — field_rename
    def test_belief_correct_field_rename(self):
        truth = {"order_field": "quantity", "required_extra": "customer_id"}
        belief = {"order_field": "quantity", "required_extra": "customer_id"}
        score = self.grader.grade_belief(belief, truth, "field_rename")
        assert score >= 0.9

    def test_belief_wrong_field_rename(self):
        truth = {"order_field": "quantity", "required_extra": "customer_id"}
        belief = {"order_field": "qty"}   # stale belief
        score = self.grader.grade_belief(belief, truth, "field_rename")
        assert score < 0.5

    # grade_belief tests — endpoint_version
    def test_belief_correct_endpoint(self):
        truth = {"rooms_endpoint": "/mock_api/v2/rooms/book"}
        belief = {"endpoint": "/mock_api/v2/rooms/book"}
        score = self.grader.grade_belief(belief, truth, "endpoint_version")
        assert score == 1.0

    def test_belief_wrong_endpoint(self):
        truth = {"rooms_endpoint": "/mock_api/v2/rooms/book"}
        belief = {"endpoint": "/mock_api/rooms/book"}
        score = self.grader.grade_belief(belief, truth, "endpoint_version")
        assert score == 0.0

    # grade_belief tests — silent_semantic
    def test_belief_correct_silent_semantic(self):
        truth = {"order_status": "approved"}
        belief = {"order_status_value": "approved"}
        score = self.grader.grade_belief(belief, truth, "silent_semantic")
        assert score == 1.0

    def test_belief_partial_silent_semantic_knows_something_changed(self):
        truth = {"order_status": "approved"}
        belief = {"drift_detected": True}   # knows something changed but not what
        score = self.grader.grade_belief(belief, truth, "silent_semantic")
        assert score == 0.4

    def test_belief_wrong_silent_semantic(self):
        truth = {"order_status": "approved"}
        belief = {"order_status_value": "confirmed"}  # stale
        score = self.grader.grade_belief(belief, truth, "silent_semantic")
        assert score == 0.0

    # grade_belief tests — policy_change
    def test_belief_correct_policy(self):
        truth = {"discount_requires_membership": True}
        belief = {"discount_requires_membership": True}
        score = self.grader.grade_belief(belief, truth, "policy_change")
        assert score == 1.0

    def test_belief_none_returns_zero(self):
        score = self.grader.grade_belief(None, {"x": 1}, "field_rename")
        assert score == 0.0

    # infer_belief_from_actions
    def test_infer_belief_probed_and_recovered(self):
        log = [
            {"step": 1, "url": "/mock_api/orders", "status": 200, "response": "{}"},
            {"step": 2, "url": "/mock_api/orders", "status": 422, "response": '{"detail": "..."}'},
            {"step": 3, "url": "/openapi.json", "status": 200, "response": "{}"},
            {"step": 4, "url": "/mock_api/orders", "status": 200, "response": '{"order_id": "x"}'},
        ]
        score = self.grader.infer_belief_from_actions(log, "field_rename")
        assert score == 0.6

    def test_infer_belief_probed_no_recovery(self):
        log = [
            {"step": 1, "url": "/mock_api/orders", "status": 200, "response": "{}"},
            {"step": 2, "url": "/mock_api/orders", "status": 422, "response": "{}"},
            {"step": 3, "url": "/openapi.json", "status": 200, "response": "{}"},
        ]
        score = self.grader.infer_belief_from_actions(log, "field_rename")
        assert score == 0.3

    def test_infer_belief_no_action(self):
        score = self.grader.infer_belief_from_actions([], "field_rename")
        assert score == 0.0


# ── 4. DriftDifficultyController ──────────────────────────────────────────────

class TestDriftDifficultyController:

    def test_initial_level_is_zero(self):
        ctrl = DriftDifficultyController()
        assert ctrl.level == 0

    def test_no_escalation_without_enough_data(self):
        ctrl = DriftDifficultyController()
        for _ in range(4):  # ESCALATION_WINDOW is 5
            ctrl.record("field_rename", 0.95)
            ctrl.record("endpoint_version", 0.95)
            ctrl.record("policy_change", 0.95)
        assert ctrl.level == 0  # not enough data yet

    def test_escalates_when_all_drift_types_above_threshold(self):
        ctrl = DriftDifficultyController()
        for _ in range(5):  # ESCALATION_WINDOW = 5
            ctrl.record("field_rename", 0.90)
            ctrl.record("endpoint_version", 0.90)
            ctrl.record("policy_change", 0.90)
        assert ctrl.level == 1

    def test_no_escalation_when_one_drift_type_below_threshold(self):
        ctrl = DriftDifficultyController()
        for _ in range(5):
            ctrl.record("field_rename", 0.90)
            ctrl.record("endpoint_version", 0.50)  # below threshold
            ctrl.record("policy_change", 0.90)
        assert ctrl.level == 0

    def test_get_scenario_params_level_zero_unchanged(self):
        ctrl = DriftDifficultyController()
        scenario = {"id": "easy_field_rename", "drift_trigger_step": 3, "drift_type": "field_rename"}
        result = ctrl.get_scenario_params(scenario)
        assert result == scenario

    def test_reset_clears_state(self):
        ctrl = DriftDifficultyController()
        ctrl._escalation_level = 2
        ctrl.reset()
        assert ctrl.level == 0
        assert len(ctrl._history) == 0


# ── 5. Reward signal verification (Day 3 Step 1 from the plan) ────────────────

class TestRewardSignal:
    """
    Manual verification that the reward signal works correctly.
    These tests don't hit a live server — they test the grader math directly.
    """

    grader = AdaptiveGrader()

    def test_untrained_behavior_produces_low_rewards(self):
        """Simulates agent that never detects drift."""
        task_reward = self.grader.grade_task(
            task_completed=False, steps_taken=8, max_steps=8, drift_detected=False
        )
        # Completely wrong belief (wrong field name; no required_extra knowledge)
        # grade_belief for field_rename scores 3 sub-checks: field_match, extra_match, ep_match
        # field_match: "qty" != "quantity" → 0.0
        # extra_match: None == "customer_id" → 0.0
        # ep_match: no claims_endpoint in truth → 1.0 (vacuous)
        # avg = (0+0+1)/3 = 0.333 — we assert it is < 0.5 (stale)
        belief_accuracy = self.grader.grade_belief(
            agent_belief={"order_field": "qty"},   # stale — wrong field name
            world_truth={"order_field": "quantity", "required_extra": "customer_id"},
            drift_type="field_rename",
        )
        combined = task_reward * 0.7 + belief_accuracy * 0.3
        assert task_reward <= 0.001
        assert belief_accuracy < 0.5   # stale belief — partial at best
        assert combined < 0.2

    def test_informed_behavior_produces_high_rewards(self):
        """Simulates agent that detects drift, probes, and adapts."""
        task_reward = self.grader.grade_task(
            task_completed=True, steps_taken=6, max_steps=8, drift_detected=True
        )
        belief_accuracy = self.grader.grade_belief(
            agent_belief={"order_field": "quantity", "required_extra": "customer_id"},
            world_truth={"order_field": "quantity", "required_extra": "customer_id"},
            drift_type="field_rename",
        )
        combined = task_reward * 0.7 + belief_accuracy * 0.3
        assert task_reward >= 0.70
        assert belief_accuracy >= 0.90
        assert combined >= 0.75

    def test_combined_reward_formula(self):
        """Verify the 0.7/0.3 split is applied correctly."""
        task_reward = 0.80
        belief_accuracy = 0.60
        expected = task_reward * 0.7 + belief_accuracy * 0.3
        # Recreate with grader outputs
        tr = self.grader.grade_task(True, 4, 8, True)
        ba = self.grader.grade_belief(
            {"order_field": "quantity", "required_extra": "customer_id"},
            {"order_field": "quantity", "required_extra": "customer_id"},
            "field_rename"
        )
        combined = tr * 0.7 + ba * 0.3
        assert 0.0 <= combined <= 1.0


# ── 6. AdaptiveState model ────────────────────────────────────────────────────

class TestModels:

    def test_adaptive_action_defaults(self):
        action = AdaptiveAction()
        assert action.action_type == "call_api"
        assert action.method == "GET"
        assert action.belief_state is None
        assert action.history_steps == 3

    def test_adaptive_observation_no_drift_occurred(self):
        """Verify drift_occurred is NOT a field in AdaptiveObservation (v2 compliance)."""
        obs = AdaptiveObservation()
        assert not hasattr(obs, "drift_occurred"), (
            "drift_occurred must be removed from AdaptiveObservation in v2"
        )
        assert not hasattr(obs, "drift_type"), (
            "drift_type must be removed from AdaptiveObservation in v2"
        )

    def test_adaptive_observation_has_v2_fields(self):
        obs = AdaptiveObservation()
        assert hasattr(obs, "prior_world_model")
        assert hasattr(obs, "episode_history")
        assert hasattr(obs, "belief_accuracy")
        assert hasattr(obs, "difficulty_level")

    def test_adaptive_state_defaults(self):
        state = AdaptiveState()
        assert not state.drift_injected
        assert state.agent_belief == {}
        assert state.world_truth == {}
        assert state.step_history == []