tests/test_smoke.py

import random

from counsel_env.models import CounselAction
from counsel_env.evaluation import evaluate_agent, present_all_agent, keyword_spam_agent, oracle_scripted_agent
from counsel_env.server.app import app
from counsel_env.server.case_generator import generate_case
from counsel_env.server.counsel_env_environment import CounselEnvironment
from counsel_env.server.witness import Contradiction, Witness


def _trigger_and_surface_first(env: CounselEnvironment):
    contradiction = env.witness.contradictions[0]
    env.step(CounselAction(tool="ask_question", text=f"{contradiction.trigger_keywords[0]}?"))
    env.step(CounselAction(tool="present_evidence", exhibit_id=contradiction.disprover_evidence_id))
    return contradiction


class TestWitnessLogic:
    def test_witness_trigger(self):
        contradiction = Contradiction(
            cid="c1",
            trigger_keywords=["where were you", "alibi"],
            sealed_claim="I was at the diner all night.",
            disprover_evidence_id="phone_log",
        )
        witness = Witness({"motive": "No motive."}, [contradiction])

        response = witness.respond_to_question("Where were you that night?")

        assert response == contradiction.sealed_claim
        assert contradiction.triggered is True
        assert witness.committed_claims == [contradiction.sealed_claim]

    def test_evidence_surface(self):
        contradiction = Contradiction(
            cid="c1",
            trigger_keywords=["alibi"],
            sealed_claim="I was at the diner.",
            disprover_evidence_id="phone_log",
        )
        witness = Witness({}, [contradiction])

        witness.respond_to_question("Tell us your alibi.")
        response = witness.react_to_evidence("phone_log")

        assert "stammers" in response.lower()
        assert contradiction.surfaced is True

    def test_no_false_positive(self):
        contradiction = Contradiction(
            cid="c1",
            trigger_keywords=["alibi"],
            sealed_claim="I was at the diner.",
            disprover_evidence_id="phone_log",
        )
        witness = Witness({}, [contradiction])

        response = witness.react_to_evidence("phone_log")

        assert "no comment" in response.lower()
        assert contradiction.surfaced is False


class TestCaseGeneration:
    def test_curriculum_difficulty_shapes_cases(self):
        easy = generate_case(difficulty="easy")
        medium = generate_case(difficulty="medium")
        hard = generate_case(difficulty="hard")

        assert len(easy["contradictions"]) == 1
        assert len(medium["contradictions"]) == 2
        assert len(hard["contradictions"]) >= 3
        assert "irrelevant_weather_report" in hard["evidence"]

    def test_curriculum_sampling_distribution(self):
        cases = [
            generate_case(curriculum_stage="mixed", distribution={"easy": 1.0})
            for _ in range(5)
        ]

        assert {case["difficulty"] for case in cases} == {"easy"}

    def test_seed_replay_is_deterministic(self):
        env_a = CounselEnvironment()
        env_b = CounselEnvironment()
        obs_a = env_a.reset(seed=1234, difficulty="hard")
        obs_b = env_b.reset(seed=1234, difficulty="hard")

        assert obs_a.case_id == obs_b.case_id
        assert obs_a.case_brief == obs_b.case_brief
        assert obs_a.evidence_descriptions == obs_b.evidence_descriptions
        assert [c.cid for c in env_a.witness.contradictions] == [c.cid for c in env_b.witness.contradictions]


class TestEnvironmentLogic:
    def test_space_demo_routes_registered(self):
        paths = {route.path for route in app.routes}

        assert "/" in paths
        assert "/demo" in paths
        assert "/demo/api/reset" in paths
        assert "/demo/api/step" in paths

    def test_reset_state(self):
        env = CounselEnvironment()
        obs = env.reset(difficulty="easy")

        assert obs.questions_remaining == 15
        assert obs.done is False
        assert obs.reward == 0.0
        assert obs.case_brief
        assert obs.available_evidence
        assert obs.evidence_descriptions
        assert set(obs.available_evidence) == set(obs.evidence_descriptions)
        assert obs.case_id == env.case["case_id"]
        assert obs.difficulty == "easy"
        assert env.state.contradictions_total == 1
        assert env.state.questions_used == 0

    def test_question_budget(self):
        env = CounselEnvironment()
        env.reset(difficulty="easy")

        obs = None
        for i in range(env.QUESTION_BUDGET):
            obs = env.step(CounselAction(tool="ask_question", text=f"Off-topic question {i}?"))

        assert obs is not None
        assert obs.done is True
        assert obs.questions_remaining == 0

        after_done = env.step(CounselAction(tool="ask_question", text="One too many?"))
        assert after_done.done is True
        assert "already over" in after_done.witness_response.lower()

    def test_duplicate_question_penalty(self):
        clean_env = CounselEnvironment()
        clean_env.reset(difficulty="easy")
        _trigger_and_surface_first(clean_env)
        clean_reward = clean_env.step(CounselAction(tool="rest_case")).reward

        duplicate_env = CounselEnvironment()
        duplicate_env.reset(difficulty="easy")
        contradiction = duplicate_env.witness.contradictions[0]
        question = f"{contradiction.trigger_keywords[0]}?"
        duplicate_env.step(CounselAction(tool="ask_question", text=question))
        duplicate_response = duplicate_env.step(CounselAction(tool="ask_question", text=question))
        duplicate_env.step(CounselAction(tool="present_evidence", exhibit_id=contradiction.disprover_evidence_id))
        duplicate_reward = duplicate_env.step(CounselAction(tool="rest_case")).reward

        assert "already asked" in duplicate_response.witness_response.lower()
        assert duplicate_env.duplicate_question_count == 1
        assert duplicate_reward < clean_reward

    def test_reward_components_primary_dominates(self):
        env = CounselEnvironment()
        env.reset(difficulty="easy")
        _trigger_and_surface_first(env)
        obs = env.step(CounselAction(tool="rest_case"))

        assert obs.reward_components["primary_reward"] == 1.0
        assert obs.reward >= 0.8
        assert obs.reward_components["auxiliary_reward_raw"] > 0.0
        assert obs.reward_components["evidence_timing_successes"] >= 1.0

    def test_no_infinite_evidence_loop(self):
        env = CounselEnvironment()
        env.reset(difficulty="easy")
        exhibit = env.available_evidence[0] if hasattr(env, "available_evidence") else env.case["evidence"].keys()
        exhibit_id = next(iter(env.case["evidence"].keys()))

        for _ in range(env.MAX_ACTIONS + 1):
            obs = env.step(CounselAction(tool="present_evidence", exhibit_id=exhibit_id))
            if obs.done:
                break

        assert obs.done is True
        assert env.action_count >= env.MAX_ACTIONS

    def test_transcript_export_contains_labels(self):
        env = CounselEnvironment()
        env.reset(seed=55, difficulty="easy", episode_id="demo")
        _trigger_and_surface_first(env)
        obs = env.step(CounselAction(tool="rest_case"))

        payload = env.export_transcript_json()
        markdown = env.export_transcript_markdown()

        assert payload["episode_id"] == "demo"
        assert payload["seed"] == 55
        assert payload["events"]
        assert payload["reward_components"]["total_reward"] == obs.reward
        assert "Triggered:" in markdown
        assert "Surfaced:" in markdown

    def test_state_tracks_extended_diagnostics(self):
        env = CounselEnvironment()
        env.reset(difficulty="easy")
        _trigger_and_surface_first(env)
        state = env.state

        assert state.contradictions_triggered == 1
        assert state.contradictions_surfaced == 1
        assert state.evidence_timing_successes == 1
        assert state.action_count >= 2


class TestRolloutValidation:
    def test_scripted_agent_success(self):
        successes = 0

        for _ in range(10):
            env = CounselEnvironment()
            env.reset(curriculum_stage="mixed")
            _trigger_and_surface_first(env)
            obs = env.step(CounselAction(tool="rest_case"))
            if obs.reward > 0:
                successes += 1

        assert successes >= 1

    def test_random_agent_baseline(self):
        rng = random.Random(7)
        rewards = []

        for _ in range(10):
            env = CounselEnvironment()
            obs = env.reset(difficulty="hard")
            for _ in range(8):
                if rng.random() < 0.7:
                    obs = env.step(
                        CounselAction(tool="ask_question", text=f"Random question {rng.randint(1, 999)}?")
                    )
                else:
                    obs = env.step(
                        CounselAction(tool="present_evidence", exhibit_id=rng.choice(obs.available_evidence))
                    )
                if obs.done:
                    break
            rewards.append(env.step(CounselAction(tool="rest_case")).reward)

        avg_reward = sum(rewards) / len(rewards)
        assert avg_reward < 0.35

    def test_reward_hacking_baselines_stay_below_oracle(self):
        seeds = [9001, 9002, 9003, 9004, 9005]
        spam_rows, _ = evaluate_agent("keyword_spam", keyword_spam_agent, seeds, transcript_limit=0)
        present_rows, _ = evaluate_agent("present_all", present_all_agent, seeds, transcript_limit=0)
        oracle_rows, _ = evaluate_agent("scripted_oracle", oracle_scripted_agent, seeds, transcript_limit=0)

        spam_reward = sum(row["reward"] for row in spam_rows) / len(spam_rows)
        present_reward = sum(row["reward"] for row in present_rows) / len(present_rows)
        oracle_reward = sum(row["reward"] for row in oracle_rows) / len(oracle_rows)

        assert oracle_reward > spam_reward
        assert oracle_reward > present_reward
        assert present_reward == 0.0


class TestEdgeCases:
    def test_empty_question(self):
        env = CounselEnvironment()
        env.reset(difficulty="easy")

        obs = env.step(CounselAction(tool="ask_question", text=""))

        assert "proper question" in obs.witness_response.lower()
        assert env.inadmissible_count == 1

    def test_invalid_exhibit(self):
        env = CounselEnvironment()
        env.reset(difficulty="easy")

        obs = env.step(CounselAction(tool="present_evidence", exhibit_id="nonexistent_exhibit"))

        assert "invalid exhibit" in obs.witness_response.lower()
        assert env.inadmissible_count == 1

    def test_extremely_long_question(self):
        env = CounselEnvironment()
        env.reset(difficulty="easy")
        long_question = "What is your alibi " + ("please explain " * 200)

        obs = env.step(CounselAction(tool="ask_question", text=long_question))

        assert obs.witness_response
        assert obs.questions_remaining == env.QUESTION_BUDGET - 1

    def test_no_trigger_keywords_case(self):
        witness = Witness(
            {"background": "I worked late."},
            [
                Contradiction(
                    cid="c1",
                    trigger_keywords=[],
                    sealed_claim="Hidden claim.",
                    disprover_evidence_id="log",
                )
            ],
        )

        response = witness.respond_to_question("What is your favorite color?")
        evidence_response = witness.react_to_evidence("log")

        assert response == "I don't recall."
        assert "no comment" in evidence_response.lower()