"""Tests for the Experiment ActiveRecord model."""

import warnings
from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock, patch

import httpx
import pytest

from lumenova_beacon.experiments.models import Experiment, _resolve_variable_mappings
from lumenova_beacon.experiments.types import (
    EnrichedRecord,
    ExperimentRun,
    ExperimentRunStats,
    ExperimentRunStatus,
    ExperimentStatus,
    ExperimentStep,
    ExperimentStepType,
)


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------

SAMPLE_EXPERIMENT_RESPONSE = {
    "id": "exp-001",
    "name": "Test Pipeline",
    "description": "A test experiment",
    "project_id": "proj-001",
    "dataset_id": "ds-001",
    "result_dataset_id": None,
    "execution_type": "server",
    "status": 1,
    "started_at": None,
    "completed_at": None,
    "created_at": "2025-01-15T10:00:00Z",
    "updated_at": None,
    "dataset_name": "Test Dataset",
    "record_count": 100,
    "step_count": 2,
    "total_runs": 0,
    "completed_runs": 0,
    "failed_runs": 0,
    "pending_runs": 0,
    "completed_records": 0,
    "failed_records": 0,
    "steps": [
        {
            "id": "step-001",
            "experiment_id": "exp-001",
            "position": 0,
            "step_type": "prompt",
            "output_column": "rendered",
            "config": {"prompt_content": "Summarize: {{input}}"},
            "created_at": "2025-01-15T10:00:00Z",
            "updated_at": None,
        },
        {
            "id": "step-002",
            "experiment_id": "exp-001",
            "position": 1,
            "step_type": "llm",
            "output_column": "response",
            "config": {"model_config_id": "cfg-001"},
            "created_at": "2025-01-15T10:00:00Z",
            "updated_at": None,
        },
    ],
}

SAMPLE_LIST_RESPONSE = {
    "items": [SAMPLE_EXPERIMENT_RESPONSE],
    "total": 1,
    "page": 1,
    "page_size": 20,
    "total_pages": 1,
}

SAMPLE_RUN_RESPONSE = {
    "id": "run-001",
    "experiment_id": "exp-001",
    "record_id": "rec-001",
    "step_id": "step-001",
    "configuration_id": None,
    "status": "completed",
    "result": {"output": "hello"},
    "error": None,
    "execution_time_ms": 150,
    "started_at": "2025-01-15T10:00:00Z",
    "completed_at": "2025-01-15T10:00:01Z",
    "created_at": "2025-01-15T10:00:00Z",
    "updated_at": "2025-01-15T10:00:01Z",
}

SAMPLE_RUNS_LIST_RESPONSE = {
    "items": [SAMPLE_RUN_RESPONSE],
    "total": 1,
    "page": 1,
    "page_size": 20,
    "total_pages": 1,
}

SAMPLE_RUN_STATS_RESPONSE = {
    "experiment_id": "exp-001",
    "total_runs": 200,
    "pending_runs": 0,
    "running_runs": 0,
    "completed_runs": 190,
    "failed_runs": 10,
    "avg_execution_time_ms": 200.5,
    "min_execution_time_ms": 50,
    "max_execution_time_ms": 3000,
}

SAMPLE_ENRICHED_RECORDS_RESPONSE = {
    "items": [
        {
            "record_id": "rec-001",
            "data": {"input": "hello"},
            "overall_status": "success",
            "step_runs": [
                {
                    "step_id": "step-001",
                    "step_position": 0,
                    "step_type": "prompt",
                    "output_column": "rendered",
                    "status": "completed",
                    "output": "Rendered text",
                },
            ],
        },
    ],
    "total": 1,
    "page": 1,
    "page_size": 20,
    "total_pages": 1,
}


def _mock_transport():
    transport = MagicMock()
    transport.headers = {"Authorization": "Bearer test"}
    transport.timeout = 30
    transport.verify = True
    return transport


def _mock_response(data, status_code=200):
    response = MagicMock(spec=httpx.Response)
    response.status_code = status_code
    response.json.return_value = data
    response.raise_for_status = MagicMock()
    return response


# ---------------------------------------------------------------------------
# Tests: _from_dict
# ---------------------------------------------------------------------------


class TestFromDict:
    def test_full_response(self):
        exp = Experiment._from_dict(SAMPLE_EXPERIMENT_RESPONSE)
        assert exp.id == "exp-001"
        assert exp.name == "Test Pipeline"
        assert exp.description == "A test experiment"
        assert exp.dataset_id == "ds-001"
        assert exp.execution_type == "server"
        assert exp.status == ExperimentStatus.DRAFT
        assert exp.record_count == 100
        assert exp.step_count == 2
        assert exp.steps is not None
        assert len(exp.steps) == 2
        assert exp.steps[0].step_type == "prompt"
        assert exp.steps[0].output_column == "rendered"
        assert exp.steps[1].step_type == "llm"
        assert exp.total_runs == 0
        assert exp.completed_records == 0
        assert exp.failed_records == 0
        assert exp.created_at == datetime(2025, 1, 15, 10, 0, tzinfo=timezone.utc)

    def test_minimal_response(self):
        data = {
            "id": "exp-002",
            "name": "Minimal",
            "dataset_id": "ds-002",
            "status": 3,
            "created_at": "2025-06-01T00:00:00Z",
        }
        exp = Experiment._from_dict(data)
        assert exp.id == "exp-002"
        assert exp.status == ExperimentStatus.RUNNING
        assert exp.steps is None
        assert exp.step_count == 0
        assert exp.execution_type is None
        assert exp.completed_records == 0
        assert exp.pending_runs == 0

    def test_completed_with_errors_status(self):
        data = {
            "id": "exp-003",
            "name": "Partial",
            "dataset_id": "ds-003",
            "status": 7,
            "created_at": "2025-06-01T00:00:00Z",
        }
        exp = Experiment._from_dict(data)
        assert exp.status == ExperimentStatus.COMPLETED_WITH_ERRORS

    def test_with_progress_counters(self):
        data = {
            **SAMPLE_EXPERIMENT_RESPONSE,
            "status": 3,
            "total_runs": 200,
            "completed_runs": 150,
            "failed_runs": 10,
            "pending_runs": 40,
            "completed_records": 70,
            "failed_records": 5,
        }
        exp = Experiment._from_dict(data)
        assert exp.total_runs == 200
        assert exp.completed_runs == 150
        assert exp.failed_runs == 10
        assert exp.pending_runs == 40
        assert exp.completed_records == 70
        assert exp.failed_records == 5


class TestRepr:
    def test_repr(self):
        exp = Experiment._from_dict(SAMPLE_EXPERIMENT_RESPONSE)
        result = repr(exp)
        assert "Test Pipeline" in result
        assert "DRAFT" in result
        assert "step_count=2" in result
        assert "record_count=100" in result


# ---------------------------------------------------------------------------
# Tests: aget / get
# ---------------------------------------------------------------------------


class TestGet:
    @pytest.mark.asyncio
    @patch("lumenova_beacon.experiments.models.get_transport")
    @patch("lumenova_beacon.experiments.models.get_base_url")
    async def test_aget(self, mock_base_url, mock_transport):
        mock_base_url.return_value = "https://api.example.com"
        mock_transport.return_value = _mock_transport()

        mock_response = _mock_response(SAMPLE_EXPERIMENT_RESPONSE)

        with patch("httpx.AsyncClient") as mock_client_cls:
            mock_client = AsyncMock()
            mock_client.get = AsyncMock(return_value=mock_response)
            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
            mock_client.__aexit__ = AsyncMock(return_value=False)
            mock_client_cls.return_value = mock_client

            exp = await Experiment.aget("exp-001")

            assert exp.id == "exp-001"
            assert exp.name == "Test Pipeline"
            mock_client.get.assert_called_once()
            call_args = mock_client.get.call_args
            assert "experiments/exp-001" in call_args[0][0]


# ---------------------------------------------------------------------------
# Tests: acreate / create
# ---------------------------------------------------------------------------


class TestCreate:
    @pytest.mark.asyncio
    @patch("lumenova_beacon.experiments.models.get_transport")
    @patch("lumenova_beacon.experiments.models.get_base_url")
    async def test_acreate_with_step_objects(self, mock_base_url, mock_transport):
        mock_base_url.return_value = "https://api.example.com"
        mock_transport.return_value = _mock_transport()

        mock_response = _mock_response(SAMPLE_EXPERIMENT_RESPONSE)

        with patch("httpx.AsyncClient") as mock_client_cls:
            mock_client = AsyncMock()
            mock_client.post = AsyncMock(return_value=mock_response)
            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
            mock_client.__aexit__ = AsyncMock(return_value=False)
            mock_client_cls.return_value = mock_client

            steps = [
                ExperimentStep(
                    step_type=ExperimentStepType.PROMPT,
                    output_column="rendered",
                    config={"prompt_content": "Hello {{name}}"},
                ),
                ExperimentStep(
                    step_type=ExperimentStepType.LLM,
                    output_column="response",
                    config={"model_config_id": "cfg-001"},
                ),
            ]

            exp = await Experiment.acreate(
                name="Test",
                dataset_id="ds-001",
                steps=steps,
                description="A test",
            )

            assert exp.id == "exp-001"
            call_args = mock_client.post.call_args
            payload = call_args[1]["json"]
            assert payload["name"] == "Test"
            assert payload["dataset_id"] == "ds-001"
            assert len(payload["steps"]) == 2
            assert payload["steps"][0]["step_type"] == "prompt"
            assert payload["steps"][1]["step_type"] == "llm"
            assert payload["execution_type"] == "server"
            assert payload["clone_dataset"] is True

    @pytest.mark.asyncio
    @patch("lumenova_beacon.experiments.models.get_transport")
    @patch("lumenova_beacon.experiments.models.get_base_url")
    async def test_acreate_with_raw_dicts(self, mock_base_url, mock_transport):
        mock_base_url.return_value = "https://api.example.com"
        mock_transport.return_value = _mock_transport()

        mock_response = _mock_response(SAMPLE_EXPERIMENT_RESPONSE)

        with patch("httpx.AsyncClient") as mock_client_cls:
            mock_client = AsyncMock()
            mock_client.post = AsyncMock(return_value=mock_response)
            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
            mock_client.__aexit__ = AsyncMock(return_value=False)
            mock_client_cls.return_value = mock_client

            steps = [
                {"step_type": "llm", "output_column": "out", "config": {}},
            ]

            exp = await Experiment.acreate(
                name="Test",
                dataset_id="ds-001",
                steps=steps,
                execution_type="local",
                clone_dataset=False,
            )

            call_args = mock_client.post.call_args
            payload = call_args[1]["json"]
            assert payload["execution_type"] == "local"
            assert payload["clone_dataset"] is False
            assert payload["steps"][0] == steps[0]


# ---------------------------------------------------------------------------
# Tests: alist / list
# ---------------------------------------------------------------------------


class TestList:
    @pytest.mark.asyncio
    @patch("lumenova_beacon.experiments.models.get_transport")
    @patch("lumenova_beacon.experiments.models.get_base_url")
    async def test_alist(self, mock_base_url, mock_transport):
        mock_base_url.return_value = "https://api.example.com"
        mock_transport.return_value = _mock_transport()

        mock_response = _mock_response(SAMPLE_LIST_RESPONSE)

        with patch("httpx.AsyncClient") as mock_client_cls:
            mock_client = AsyncMock()
            mock_client.get = AsyncMock(return_value=mock_response)
            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
            mock_client.__aexit__ = AsyncMock(return_value=False)
            mock_client_cls.return_value = mock_client

            experiments, pagination = await Experiment.alist(
                page=1, page_size=20, status=ExperimentStatus.DRAFT
            )

            assert len(experiments) == 1
            assert experiments[0].id == "exp-001"
            assert pagination.total == 1
            assert pagination.total_pages == 1


# ---------------------------------------------------------------------------
# Tests: aclone
# ---------------------------------------------------------------------------


class TestClone:
    @pytest.mark.asyncio
    @patch("lumenova_beacon.experiments.models.get_transport")
    @patch("lumenova_beacon.experiments.models.get_base_url")
    async def test_aclone(self, mock_base_url, mock_transport):
        mock_base_url.return_value = "https://api.example.com"
        mock_transport.return_value = _mock_transport()

        cloned_response = {**SAMPLE_EXPERIMENT_RESPONSE, "id": "exp-cloned"}
        mock_response = _mock_response(cloned_response)

        with patch("httpx.AsyncClient") as mock_client_cls:
            mock_client = AsyncMock()
            mock_client.post = AsyncMock(return_value=mock_response)
            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
            mock_client.__aexit__ = AsyncMock(return_value=False)
            mock_client_cls.return_value = mock_client

            exp = Experiment._from_dict(SAMPLE_EXPERIMENT_RESPONSE)
            cloned = await exp.aclone()

            assert cloned.id == "exp-cloned"
            assert cloned.id != exp.id
            call_args = mock_client.post.call_args
            assert "clone" in call_args[0][0]


# ---------------------------------------------------------------------------
# Tests: areset / reset
# ---------------------------------------------------------------------------


class TestReset:
    @pytest.mark.asyncio
    @patch("lumenova_beacon.experiments.models.get_transport")
    @patch("lumenova_beacon.experiments.models.get_base_url")
    async def test_areset(self, mock_base_url, mock_transport):
        mock_base_url.return_value = "https://api.example.com"
        mock_transport.return_value = _mock_transport()

        reset_response = {
            **SAMPLE_EXPERIMENT_RESPONSE,
            "status": 1,  # DRAFT
            "started_at": None,
            "completed_at": None,
            "total_runs": 0,
            "completed_runs": 0,
            "failed_runs": 0,
            "pending_runs": 0,
            "completed_records": 0,
            "failed_records": 0,
        }
        mock_response = _mock_response(reset_response)

        with patch("httpx.AsyncClient") as mock_client_cls:
            mock_client = AsyncMock()
            mock_client.post = AsyncMock(return_value=mock_response)
            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
            mock_client.__aexit__ = AsyncMock(return_value=False)
            mock_client_cls.return_value = mock_client

            exp = Experiment._from_dict(
                {**SAMPLE_EXPERIMENT_RESPONSE, "status": 4, "completed_runs": 100}
            )
            assert exp.status == ExperimentStatus.COMPLETED

            result = await exp.areset()

            assert result.status == ExperimentStatus.DRAFT
            assert result.total_runs == 0
            assert result.completed_runs == 0
            call_args = mock_client.post.call_args
            assert "reset" in call_args[0][0]


# ---------------------------------------------------------------------------
# Tests: acontinue_experiment / continue_experiment
# ---------------------------------------------------------------------------


class TestContinueExperiment:
    @pytest.mark.asyncio
    @patch("lumenova_beacon.experiments.models.get_transport")
    @patch("lumenova_beacon.experiments.models.get_base_url")
    async def test_acontinue_stopped(self, mock_base_url, mock_transport):
        """STOPPED experiment POSTs to /continue endpoint."""
        mock_base_url.return_value = "https://api.example.com"
        mock_transport.return_value = _mock_transport()

        continued_response = {
            **SAMPLE_EXPERIMENT_RESPONSE,
            "status": 2,  # QUEUED
            "total_runs": 200,
            "completed_runs": 100,
            "pending_runs": 100,
        }
        mock_response = _mock_response(continued_response)

        with patch("httpx.AsyncClient") as mock_client_cls:
            mock_client = AsyncMock()
            mock_client.post = AsyncMock(return_value=mock_response)
            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
            mock_client.__aexit__ = AsyncMock(return_value=False)
            mock_client_cls.return_value = mock_client

            exp = Experiment._from_dict(
                {**SAMPLE_EXPERIMENT_RESPONSE, "status": 6, "completed_runs": 100}
            )
            assert exp.status == ExperimentStatus.STOPPED

            result = await exp.acontinue_experiment()

            assert result.status == ExperimentStatus.QUEUED
            assert result.completed_runs == 100
            call_args = mock_client.post.call_args
            assert "continue" in call_args[0][0]

    @pytest.mark.asyncio
    async def test_acontinue_waiting_for_external_delegates_to_astart(self):
        """WAITING_FOR_EXTERNAL skips POST /continue and delegates to astart."""
        exp = Experiment._from_dict(
            {**SAMPLE_EXPERIMENT_RESPONSE, "status": 8}  # WAITING_FOR_EXTERNAL
        )
        assert exp.status == ExperimentStatus.WAITING_FOR_EXTERNAL

        with patch.object(
            exp, "astart", new_callable=AsyncMock, return_value=exp
        ) as mock_astart:
            result = await exp.acontinue_experiment(
                external_agents={"My Agent": lambda x: x},
                poll_interval=3.0,
                trace_calls=False,
            )

            mock_astart.assert_called_once_with(
                external_agents={"My Agent": lambda x: x},
                poll_interval=3.0,
                trace_calls=False,
            )
            assert result is exp

    @pytest.mark.asyncio
    @patch("lumenova_beacon.experiments.models.get_transport")
    @patch("lumenova_beacon.experiments.models.get_base_url")
    async def test_acontinue_stopped_with_external_agents(
        self, mock_base_url, mock_transport
    ):
        """STOPPED experiment POSTs /continue then enters astart polling loop."""
        mock_base_url.return_value = "https://api.example.com"
        mock_transport.return_value = _mock_transport()

        continued_response = {
            **SAMPLE_EXPERIMENT_RESPONSE,
            "status": 2,  # QUEUED
        }
        mock_response = _mock_response(continued_response)

        with patch("httpx.AsyncClient") as mock_client_cls:
            mock_client = AsyncMock()
            mock_client.post = AsyncMock(return_value=mock_response)
            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
            mock_client.__aexit__ = AsyncMock(return_value=False)
            mock_client_cls.return_value = mock_client

            exp = Experiment._from_dict(
                {**SAMPLE_EXPERIMENT_RESPONSE, "status": 6}  # STOPPED
            )

            with patch.object(
                exp, "astart", new_callable=AsyncMock, return_value=exp
            ) as mock_astart:
                result = await exp.acontinue_experiment(
                    external_agents={"My Agent": lambda x: x},
                )

                # _post_continue was called (POST to /continue)
                mock_client.post.assert_called_once()
                assert "continue" in mock_client.post.call_args[0][0]

                # astart was called with external_agents for polling
                mock_astart.assert_called_once_with(
                    external_agents={"My Agent": lambda x: x},
                    poll_interval=2.0,
                    trace_calls=True,
                )


# ---------------------------------------------------------------------------
# Tests: apreview_output / preview_output
# ---------------------------------------------------------------------------


class TestPreviewOutput:
    @pytest.mark.asyncio
    @patch("lumenova_beacon.experiments.models.get_transport")
    @patch("lumenova_beacon.experiments.models.get_base_url")
    async def test_apreview_output(self, mock_base_url, mock_transport):
        mock_base_url.return_value = "https://api.example.com"
        mock_transport.return_value = _mock_transport()

        preview_response = {
            "results": [
                {
                    "configuration_label": "A",
                    "record_id": "rec-001",
                    "status": "completed",
                    "output": "Preview output",
                    "error_message": None,
                    "execution_time_ms": 200,
                    "metadata": {"model": "gpt-4"},
                },
            ],
            "total_execution_time_ms": 200,
        }
        mock_response = _mock_response(preview_response)

        with patch("httpx.AsyncClient") as mock_client_cls:
            mock_client = AsyncMock()
            mock_client.post = AsyncMock(return_value=mock_response)
            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
            mock_client.__aexit__ = AsyncMock(return_value=False)
            mock_client_cls.return_value = mock_client

            result = await Experiment.apreview_output(
                dataset_id="ds-001",
                configurations=[
                    {
                        "label": "A",
                        "prompt_content": "Hello {{name}}",
                        "model_config_id": "cfg-001",
                        "model_parameters": {"temperature": 0.7},
                    },
                ],
                record_ids=["rec-001"],
                variable_mappings={"name": "user_name"},
            )

            assert result["total_execution_time_ms"] == 200
            assert len(result["results"]) == 1
            assert result["results"][0]["output"] == "Preview output"
            call_args = mock_client.post.call_args
            assert "preview-output" in call_args[0][0]
            payload = call_args[1]["json"]
            assert payload["dataset_id"] == "ds-001"
            assert payload["record_ids"] == ["rec-001"]
            assert payload["variable_mappings"] == {"name": "user_name"}

    @pytest.mark.asyncio
    @patch("lumenova_beacon.experiments.models.get_transport")
    @patch("lumenova_beacon.experiments.models.get_base_url")
    async def test_apreview_output_without_variable_mappings(
        self, mock_base_url, mock_transport
    ):
        mock_base_url.return_value = "https://api.example.com"
        mock_transport.return_value = _mock_transport()

        preview_response = {"results": [], "total_execution_time_ms": 0}
        mock_response = _mock_response(preview_response)

        with patch("httpx.AsyncClient") as mock_client_cls:
            mock_client = AsyncMock()
            mock_client.post = AsyncMock(return_value=mock_response)
            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
            mock_client.__aexit__ = AsyncMock(return_value=False)
            mock_client_cls.return_value = mock_client

            await Experiment.apreview_output(
                dataset_id="ds-001",
                configurations=[{"label": "A"}],
                record_ids=["rec-001"],
            )

            payload = mock_client.post.call_args[1]["json"]
            assert "variable_mappings" not in payload


# ---------------------------------------------------------------------------
# Tests: alist_runs / get_run / get_run_stats
# ---------------------------------------------------------------------------


class TestRuns:
    @pytest.mark.asyncio
    @patch("lumenova_beacon.experiments.models.get_transport")
    @patch("lumenova_beacon.experiments.models.get_base_url")
    async def test_alist_runs(self, mock_base_url, mock_transport):
        mock_base_url.return_value = "https://api.example.com"
        mock_transport.return_value = _mock_transport()

        mock_response = _mock_response(SAMPLE_RUNS_LIST_RESPONSE)

        with patch("httpx.AsyncClient") as mock_client_cls:
            mock_client = AsyncMock()
            mock_client.get = AsyncMock(return_value=mock_response)
            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
            mock_client.__aexit__ = AsyncMock(return_value=False)
            mock_client_cls.return_value = mock_client

            exp = Experiment._from_dict(SAMPLE_EXPERIMENT_RESPONSE)
            runs, pagination = await exp.alist_runs(
                status=ExperimentRunStatus.COMPLETED
            )

            assert len(runs) == 1
            assert isinstance(runs[0], ExperimentRun)
            assert runs[0].status == ExperimentRunStatus.COMPLETED
            assert pagination.total == 1

    @pytest.mark.asyncio
    @patch("lumenova_beacon.experiments.models.get_transport")
    @patch("lumenova_beacon.experiments.models.get_base_url")
    async def test_aget_run(self, mock_base_url, mock_transport):
        mock_base_url.return_value = "https://api.example.com"
        mock_transport.return_value = _mock_transport()

        mock_response = _mock_response(SAMPLE_RUN_RESPONSE)

        with patch("httpx.AsyncClient") as mock_client_cls:
            mock_client = AsyncMock()
            mock_client.get = AsyncMock(return_value=mock_response)
            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
            mock_client.__aexit__ = AsyncMock(return_value=False)
            mock_client_cls.return_value = mock_client

            exp = Experiment._from_dict(SAMPLE_EXPERIMENT_RESPONSE)
            run = await exp.aget_run("run-001")

            assert isinstance(run, ExperimentRun)
            assert run.id == "run-001"
            assert run.execution_time_ms == 150
            call_args = mock_client.get.call_args
            assert "runs/run-001" in call_args[0][0]

    @pytest.mark.asyncio
    @patch("lumenova_beacon.experiments.models.get_transport")
    @patch("lumenova_beacon.experiments.models.get_base_url")
    async def test_aget_run_stats(self, mock_base_url, mock_transport):
        mock_base_url.return_value = "https://api.example.com"
        mock_transport.return_value = _mock_transport()

        mock_response = _mock_response(SAMPLE_RUN_STATS_RESPONSE)

        with patch("httpx.AsyncClient") as mock_client_cls:
            mock_client = AsyncMock()
            mock_client.get = AsyncMock(return_value=mock_response)
            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
            mock_client.__aexit__ = AsyncMock(return_value=False)
            mock_client_cls.return_value = mock_client

            exp = Experiment._from_dict(SAMPLE_EXPERIMENT_RESPONSE)
            stats = await exp.aget_run_stats()

            assert isinstance(stats, ExperimentRunStats)
            assert stats.total_runs == 200
            assert stats.completed_runs == 190
            assert stats.failed_runs == 10
            assert stats.avg_execution_time_ms == 200.5


# ---------------------------------------------------------------------------
# Tests: alist_records
# ---------------------------------------------------------------------------


class TestRecords:
    @pytest.mark.asyncio
    @patch("lumenova_beacon.experiments.models.get_transport")
    @patch("lumenova_beacon.experiments.models.get_base_url")
    async def test_alist_records(self, mock_base_url, mock_transport):
        mock_base_url.return_value = "https://api.example.com"
        mock_transport.return_value = _mock_transport()

        mock_response = _mock_response(SAMPLE_ENRICHED_RECORDS_RESPONSE)

        with patch("httpx.AsyncClient") as mock_client_cls:
            mock_client = AsyncMock()
            mock_client.get = AsyncMock(return_value=mock_response)
            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
            mock_client.__aexit__ = AsyncMock(return_value=False)
            mock_client_cls.return_value = mock_client

            exp = Experiment._from_dict(SAMPLE_EXPERIMENT_RESPONSE)
            records, pagination = await exp.alist_records(status="success")

            assert len(records) == 1
            assert isinstance(records[0], EnrichedRecord)
            assert records[0].overall_status == "success"
            assert len(records[0].step_runs) == 1
            assert records[0].step_runs[0].step_type == "prompt"


# ---------------------------------------------------------------------------
# Tests: aretry_failed / abulk_rerun
# ---------------------------------------------------------------------------


class TestRetryAndRerun:
    @pytest.mark.asyncio
    @patch("lumenova_beacon.experiments.models.get_transport")
    @patch("lumenova_beacon.experiments.models.get_base_url")
    async def test_aretry_failed(self, mock_base_url, mock_transport):
        mock_base_url.return_value = "https://api.example.com"
        mock_transport.return_value = _mock_transport()

        retry_response = {
            "retried_count": 5,
            "experiment_id": "exp-001",
            "retried_run_ids": ["r1", "r2", "r3", "r4", "r5"],
        }
        mock_response = _mock_response(retry_response)

        with patch("httpx.AsyncClient") as mock_client_cls:
            mock_client = AsyncMock()
            mock_client.post = AsyncMock(return_value=mock_response)
            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
            mock_client.__aexit__ = AsyncMock(return_value=False)
            mock_client_cls.return_value = mock_client

            exp = Experiment._from_dict(SAMPLE_EXPERIMENT_RESPONSE)
            result = await exp.aretry_failed()

            assert result["retried_count"] == 5
            call_args = mock_client.post.call_args
            assert "retry-failed" in call_args[0][0]

    @pytest.mark.asyncio
    @patch("lumenova_beacon.experiments.models.get_transport")
    @patch("lumenova_beacon.experiments.models.get_base_url")
    async def test_abulk_rerun(self, mock_base_url, mock_transport):
        mock_base_url.return_value = "https://api.example.com"
        mock_transport.return_value = _mock_transport()

        rerun_response = {
            "retried_count": 2,
            "experiment_id": "exp-001",
            "retried_run_ids": ["r1", "r2"],
        }
        mock_response = _mock_response(rerun_response)

        with patch("httpx.AsyncClient") as mock_client_cls:
            mock_client = AsyncMock()
            mock_client.post = AsyncMock(return_value=mock_response)
            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
            mock_client.__aexit__ = AsyncMock(return_value=False)
            mock_client_cls.return_value = mock_client

            exp = Experiment._from_dict(SAMPLE_EXPERIMENT_RESPONSE)
            result = await exp.abulk_rerun(["r1", "r2"])

            assert result["retried_count"] == 2
            call_args = mock_client.post.call_args
            assert "bulk-rerun" in call_args[0][0]
            assert call_args[1]["json"] == {"run_ids": ["r1", "r2"]}


# ---------------------------------------------------------------------------
# Tests: aupdate_status / aupdate_progress / aupdate_run
# ---------------------------------------------------------------------------


class TestSDKManaged:
    @pytest.mark.asyncio
    @patch("lumenova_beacon.experiments.models.get_transport")
    @patch("lumenova_beacon.experiments.models.get_base_url")
    async def test_aupdate_status(self, mock_base_url, mock_transport):
        mock_base_url.return_value = "https://api.example.com"
        mock_transport.return_value = _mock_transport()

        updated_response = {**SAMPLE_EXPERIMENT_RESPONSE, "status": 3}
        mock_response = _mock_response(updated_response)

        with patch("httpx.AsyncClient") as mock_client_cls:
            mock_client = AsyncMock()
            mock_client.patch = AsyncMock(return_value=mock_response)
            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
            mock_client.__aexit__ = AsyncMock(return_value=False)
            mock_client_cls.return_value = mock_client

            exp = Experiment._from_dict(SAMPLE_EXPERIMENT_RESPONSE)
            result = await exp.aupdate_status(ExperimentStatus.RUNNING)

            assert result.status == ExperimentStatus.RUNNING
            call_args = mock_client.patch.call_args
            assert "status" in call_args[0][0]
            assert call_args[1]["json"] == {"status": 3}

    @pytest.mark.asyncio
    @patch("lumenova_beacon.experiments.models.get_transport")
    @patch("lumenova_beacon.experiments.models.get_base_url")
    async def test_aupdate_progress(self, mock_base_url, mock_transport):
        mock_base_url.return_value = "https://api.example.com"
        mock_transport.return_value = _mock_transport()

        updated_response = {
            **SAMPLE_EXPERIMENT_RESPONSE,
            "total_runs": 200,
            "completed_runs": 100,
            "failed_runs": 5,
        }
        mock_response = _mock_response(updated_response)

        with patch("httpx.AsyncClient") as mock_client_cls:
            mock_client = AsyncMock()
            mock_client.patch = AsyncMock(return_value=mock_response)
            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
            mock_client.__aexit__ = AsyncMock(return_value=False)
            mock_client_cls.return_value = mock_client

            exp = Experiment._from_dict(SAMPLE_EXPERIMENT_RESPONSE)
            result = await exp.aupdate_progress(
                total_runs=200, completed_runs=100, failed_runs=5
            )

            assert result.total_runs == 200
            assert result.completed_runs == 100
            call_args = mock_client.patch.call_args
            assert call_args[1]["json"] == {
                "total_runs": 200,
                "completed_runs": 100,
                "failed_runs": 5,
            }

    @pytest.mark.asyncio
    @patch("lumenova_beacon.experiments.models.get_transport")
    @patch("lumenova_beacon.experiments.models.get_base_url")
    async def test_aupdate_run(self, mock_base_url, mock_transport):
        mock_base_url.return_value = "https://api.example.com"
        mock_transport.return_value = _mock_transport()

        mock_response = _mock_response(SAMPLE_RUN_RESPONSE)

        with patch("httpx.AsyncClient") as mock_client_cls:
            mock_client = AsyncMock()
            mock_client.patch = AsyncMock(return_value=mock_response)
            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
            mock_client.__aexit__ = AsyncMock(return_value=False)
            mock_client_cls.return_value = mock_client

            exp = Experiment._from_dict(SAMPLE_EXPERIMENT_RESPONSE)
            run = await exp.aupdate_run(
                "run-001",
                output="Agent result",
                status="completed",
                execution_time_ms=500,
            )

            assert isinstance(run, ExperimentRun)
            call_args = mock_client.patch.call_args
            assert "runs/run-001" in call_args[0][0]
            payload = call_args[1]["json"]
            assert payload["output"] == "Agent result"
            assert payload["status"] == "completed"
            assert payload["execution_time_ms"] == 500

    @pytest.mark.asyncio
    @patch("lumenova_beacon.experiments.models.get_transport")
    @patch("lumenova_beacon.experiments.models.get_base_url")
    async def test_aupdate_run_with_execution_metadata(
        self, mock_base_url, mock_transport
    ):
        mock_base_url.return_value = "https://api.example.com"
        mock_transport.return_value = _mock_transport()

        mock_response = _mock_response(SAMPLE_RUN_RESPONSE)

        with patch("httpx.AsyncClient") as mock_client_cls:
            mock_client = AsyncMock()
            mock_client.patch = AsyncMock(return_value=mock_response)
            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
            mock_client.__aexit__ = AsyncMock(return_value=False)
            mock_client_cls.return_value = mock_client

            exp = Experiment._from_dict(SAMPLE_EXPERIMENT_RESPONSE)
            metadata = {"model": "gpt-4", "tokens": 150}
            await exp.aupdate_run(
                "run-001",
                output="Result",
                execution_metadata=metadata,
            )

            payload = mock_client.patch.call_args[1]["json"]
            assert payload["execution_metadata"] == metadata

    @pytest.mark.asyncio
    @patch("lumenova_beacon.experiments.models.get_transport")
    @patch("lumenova_beacon.experiments.models.get_base_url")
    async def test_aupdate_run_without_execution_metadata(
        self, mock_base_url, mock_transport
    ):
        mock_base_url.return_value = "https://api.example.com"
        mock_transport.return_value = _mock_transport()

        mock_response = _mock_response(SAMPLE_RUN_RESPONSE)

        with patch("httpx.AsyncClient") as mock_client_cls:
            mock_client = AsyncMock()
            mock_client.patch = AsyncMock(return_value=mock_response)
            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
            mock_client.__aexit__ = AsyncMock(return_value=False)
            mock_client_cls.return_value = mock_client

            exp = Experiment._from_dict(SAMPLE_EXPERIMENT_RESPONSE)
            await exp.aupdate_run("run-001", output="Result")

            payload = mock_client.patch.call_args[1]["json"]
            assert "execution_metadata" not in payload


# ---------------------------------------------------------------------------
# Tests: aupdate_draft
# ---------------------------------------------------------------------------


class TestUpdateDraft:
    @pytest.mark.asyncio
    @patch("lumenova_beacon.experiments.models.get_transport")
    @patch("lumenova_beacon.experiments.models.get_base_url")
    async def test_aupdate_draft(self, mock_base_url, mock_transport):
        mock_base_url.return_value = "https://api.example.com"
        mock_transport.return_value = _mock_transport()

        updated_response = {
            **SAMPLE_EXPERIMENT_RESPONSE,
            "name": "Updated Name",
            "steps": [
                {
                    "id": "step-new",
                    "experiment_id": "exp-001",
                    "position": 0,
                    "step_type": "external_agent",
                    "output_column": "agent_out",
                    "config": {},
                    "created_at": "2025-01-15T10:00:00Z",
                    "updated_at": None,
                },
            ],
        }
        mock_response = _mock_response(updated_response)

        with patch("httpx.AsyncClient") as mock_client_cls:
            mock_client = AsyncMock()
            mock_client.put = AsyncMock(return_value=mock_response)
            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
            mock_client.__aexit__ = AsyncMock(return_value=False)
            mock_client_cls.return_value = mock_client

            exp = Experiment._from_dict(SAMPLE_EXPERIMENT_RESPONSE)
            new_steps = [
                ExperimentStep(
                    step_type=ExperimentStepType.EXTERNAL_AGENT,
                    output_column="agent_out",
                    config={},
                ),
            ]
            result = await exp.aupdate_draft(
                name="Updated Name", steps=new_steps
            )

            assert result.name == "Updated Name"
            assert result.steps is not None
            assert len(result.steps) == 1
            assert result.steps[0].step_type == "external_agent"
            call_args = mock_client.put.call_args
            assert "draft" in call_args[0][0]

    @pytest.mark.asyncio
    @patch("lumenova_beacon.experiments.models.get_transport")
    @patch("lumenova_beacon.experiments.models.get_base_url")
    async def test_aupdate_draft_does_not_send_execution_type(
        self, mock_base_url, mock_transport
    ):
        """execution_type is not part of the server's ExperimentDraftUpdate schema."""
        mock_base_url.return_value = "https://api.example.com"
        mock_transport.return_value = _mock_transport()

        mock_response = _mock_response(SAMPLE_EXPERIMENT_RESPONSE)

        with patch("httpx.AsyncClient") as mock_client_cls:
            mock_client = AsyncMock()
            mock_client.put = AsyncMock(return_value=mock_response)
            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
            mock_client.__aexit__ = AsyncMock(return_value=False)
            mock_client_cls.return_value = mock_client

            exp = Experiment._from_dict(SAMPLE_EXPERIMENT_RESPONSE)
            await exp.aupdate_draft(name="New Name")

            payload = mock_client.put.call_args[1]["json"]
            assert "execution_type" not in payload


# ---------------------------------------------------------------------------
# Tests: _update_from_response helpers
# ---------------------------------------------------------------------------


class TestUpdateHelpers:
    def test_update_from_response(self):
        exp = Experiment._from_dict(SAMPLE_EXPERIMENT_RESPONSE)
        assert exp.status == ExperimentStatus.DRAFT

        exp._update_from_response({"status": 4, "completed_runs": 100})
        assert exp.status == ExperimentStatus.COMPLETED
        assert exp.completed_runs == 100
        # Fields not in response should keep old values
        assert exp.name == "Test Pipeline"

    def test_update_from_response_obj(self):
        exp1 = Experiment._from_dict(SAMPLE_EXPERIMENT_RESPONSE)
        exp2 = Experiment._from_dict(
            {**SAMPLE_EXPERIMENT_RESPONSE, "name": "Updated", "status": 4}
        )

        exp1._update_from_response_obj(exp2)
        assert exp1.name == "Updated"
        assert exp1.status == ExperimentStatus.COMPLETED


# ---------------------------------------------------------------------------
# Tests: _warn_if_has_external_steps
# ---------------------------------------------------------------------------


class TestWarnIfHasExternalSteps:
    def test_warns_when_external_steps_present(self):
        exp = Experiment._from_dict(SAMPLE_EXTERNAL_AGENT_EXPERIMENT)
        assert exp.steps is not None

        with pytest.warns(UserWarning, match="My Agent"):
            exp._warn_if_has_external_steps()

    def test_no_warning_without_external_steps(self):
        exp = Experiment._from_dict(SAMPLE_EXPERIMENT_RESPONSE)
        assert exp.steps is not None

        with warnings.catch_warnings():
            warnings.simplefilter("error")
            exp._warn_if_has_external_steps()

    def test_no_warning_when_steps_not_loaded(self):
        exp = Experiment._from_dict(
            {**SAMPLE_EXPERIMENT_RESPONSE, "steps": None}
        )
        assert exp.steps is None

        with warnings.catch_warnings():
            warnings.simplefilter("error")
            exp._warn_if_has_external_steps()

    @pytest.mark.asyncio
    @patch("lumenova_beacon.experiments.models.get_transport")
    @patch("lumenova_beacon.experiments.models.get_base_url")
    async def test_astart_fire_and_forget_warns_with_external_steps(
        self, mock_base_url, mock_transport
    ):
        """astart() without external_agents warns when experiment has external steps."""
        mock_base_url.return_value = "https://api.example.com"
        mock_transport.return_value = _mock_transport()

        started_response = {**SAMPLE_EXTERNAL_AGENT_EXPERIMENT, "status": 2}
        mock_response = _mock_response(started_response)

        with patch("httpx.AsyncClient") as mock_client_cls:
            mock_client = AsyncMock()
            mock_client.post = AsyncMock(return_value=mock_response)
            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
            mock_client.__aexit__ = AsyncMock(return_value=False)
            mock_client_cls.return_value = mock_client

            exp = Experiment._from_dict(
                {**SAMPLE_EXTERNAL_AGENT_EXPERIMENT, "status": 1}  # DRAFT
            )

            with pytest.warns(UserWarning, match="My Agent"):
                await exp.astart()  # No external_agents

    @pytest.mark.asyncio
    @patch("lumenova_beacon.experiments.models.get_transport")
    @patch("lumenova_beacon.experiments.models.get_base_url")
    async def test_acontinue_stopped_warns_without_external_agents(
        self, mock_base_url, mock_transport
    ):
        """continue_experiment() on STOPPED without agents warns about external steps."""
        mock_base_url.return_value = "https://api.example.com"
        mock_transport.return_value = _mock_transport()

        continued_response = {**SAMPLE_EXTERNAL_AGENT_EXPERIMENT, "status": 2}
        mock_response = _mock_response(continued_response)

        with patch("httpx.AsyncClient") as mock_client_cls:
            mock_client = AsyncMock()
            mock_client.post = AsyncMock(return_value=mock_response)
            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
            mock_client.__aexit__ = AsyncMock(return_value=False)
            mock_client_cls.return_value = mock_client

            exp = Experiment._from_dict(
                {**SAMPLE_EXTERNAL_AGENT_EXPERIMENT, "status": 6}  # STOPPED
            )

            with pytest.warns(UserWarning, match="My Agent"):
                await exp.acontinue_experiment()  # No external_agents


# ---------------------------------------------------------------------------
# Tests: _resolve_variable_mappings
# ---------------------------------------------------------------------------


class TestResolveVariableMappings:
    def test_resolve_column_reference(self):
        config = {"variable_mappings": {"input": "text_column"}}
        record_data = {"text_column": "Hello world"}
        result = _resolve_variable_mappings(config, record_data)
        assert result == {"input": "Hello world"}

    def test_resolve_fixed_value(self):
        config = {
            "variable_mappings": {
                "input": {"is_fixed": True, "fixed_value": "constant text"},
            }
        }
        record_data = {"text_column": "ignored"}
        result = _resolve_variable_mappings(config, record_data)
        assert result == {"input": "constant text"}

    def test_resolve_missing_column(self):
        config = {"variable_mappings": {"input": "nonexistent"}}
        record_data = {"text_column": "Hello"}
        result = _resolve_variable_mappings(config, record_data)
        assert result == {"input": ""}

    def test_resolve_multiple_mappings(self):
        config = {
            "variable_mappings": {
                "input": "text",
                "context": "ctx_column",
            }
        }
        record_data = {"text": "query", "ctx_column": "background info"}
        result = _resolve_variable_mappings(config, record_data)
        assert result == {"input": "query", "context": "background info"}

    def test_resolve_empty_mappings(self):
        config = {}
        record_data = {"text": "hello"}
        result = _resolve_variable_mappings(config, record_data)
        assert result == {}

    def test_resolve_non_string_value(self):
        config = {"variable_mappings": {"input": "count"}}
        record_data = {"count": 42}
        result = _resolve_variable_mappings(config, record_data)
        assert result == {"input": "42"}


# ---------------------------------------------------------------------------
# Tests: astart with external_agents
# ---------------------------------------------------------------------------

# Experiment response with external_agent step (WAITING_FOR_EXTERNAL status)
SAMPLE_EXTERNAL_AGENT_EXPERIMENT = {
    **SAMPLE_EXPERIMENT_RESPONSE,
    "status": 8,  # WAITING_FOR_EXTERNAL
    "steps": [
        {
            "id": "step-001",
            "experiment_id": "exp-001",
            "position": 0,
            "step_type": "external_agent",
            "output_column": "agent_output",
            "config": {
                "agent_name": "My Agent",
                "variable_mappings": {"input": "text"},
            },
            "created_at": "2025-01-15T10:00:00Z",
            "updated_at": None,
        },
    ],
}

SAMPLE_PENDING_RUN = {
    "id": "run-001",
    "experiment_id": "exp-001",
    "record_id": "rec-001",
    "step_id": "step-001",
    "configuration_id": None,
    "status": "pending",
    "result": None,
    "error": None,
    "execution_time_ms": None,
    "started_at": None,
    "completed_at": None,
    "created_at": "2025-01-15T10:00:00Z",
    "updated_at": None,
}


class TestStartWithExternalAgents:
    @pytest.mark.asyncio
    @patch("lumenova_beacon.experiments.models.get_transport")
    @patch("lumenova_beacon.experiments.models.get_base_url")
    async def test_astart_fire_and_forget(self, mock_base_url, mock_transport):
        """start() with no external_agents just POSTs to /start and returns."""
        mock_base_url.return_value = "https://api.example.com"
        mock_transport.return_value = _mock_transport()

        started_response = {**SAMPLE_EXPERIMENT_RESPONSE, "status": 2}
        mock_response = _mock_response(started_response)

        with patch("httpx.AsyncClient") as mock_client_cls:
            mock_client = AsyncMock()
            mock_client.post = AsyncMock(return_value=mock_response)
            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
            mock_client.__aexit__ = AsyncMock(return_value=False)
            mock_client_cls.return_value = mock_client

            exp = Experiment._from_dict(SAMPLE_EXPERIMENT_RESPONSE)
            assert exp.status == ExperimentStatus.DRAFT

            result = await exp.astart()

            assert result.status == ExperimentStatus.QUEUED
            mock_client.post.assert_called_once()
            call_args = mock_client.post.call_args
            assert "start" in call_args[0][0]

    @pytest.mark.asyncio
    async def test_astart_returns_on_terminal_status(self):
        """Polling loop exits when experiment reaches terminal status."""
        exp = Experiment._from_dict(
            {**SAMPLE_EXTERNAL_AGENT_EXPERIMENT, "status": 3}  # RUNNING
        )

        # Mock aget to return COMPLETED on first poll
        completed_exp = Experiment._from_dict(
            {**SAMPLE_EXTERNAL_AGENT_EXPERIMENT, "status": 4}  # COMPLETED
        )

        with patch.object(Experiment, "aget", new_callable=AsyncMock) as mock_aget:
            mock_aget.return_value = completed_exp

            result = await exp.astart(
                external_agents={"My Agent": lambda x: x}
            )

            assert result.status == ExperimentStatus.COMPLETED

    @pytest.mark.asyncio
    async def test_astart_returns_when_no_matching_agents(self):
        """Returns when pending external runs exist but none match provided agents."""
        exp = Experiment._from_dict(SAMPLE_EXTERNAL_AGENT_EXPERIMENT)

        # Mock aget to return WAITING_FOR_EXTERNAL
        with patch.object(Experiment, "aget", new_callable=AsyncMock) as mock_aget:
            mock_aget.return_value = Experiment._from_dict(
                SAMPLE_EXTERNAL_AGENT_EXPERIMENT
            )

            # Mock pending runs — step has agent_name="My Agent" but we provide "Other Agent"
            pending_runs_response = {
                "items": [SAMPLE_PENDING_RUN],
                "total": 1,
                "page": 1,
                "page_size": 100,
                "total_pages": 1,
            }
            with patch.object(
                exp,
                "alist_runs",
                new_callable=AsyncMock,
                return_value=(
                    [ExperimentRun.from_dict(SAMPLE_PENDING_RUN)],
                    MagicMock(total_pages=1),
                ),
            ):
                result = await exp.astart(
                    external_agents={"Other Agent": lambda x: x}
                )

                # Should return because no runs match "Other Agent"
                assert result.status == ExperimentStatus.WAITING_FOR_EXTERNAL

    @pytest.mark.asyncio
    async def test_astart_processes_matching_runs(self):
        """Executes matching external agent runs and continues polling."""
        exp = Experiment._from_dict(SAMPLE_EXTERNAL_AGENT_EXPERIMENT)

        call_count = 0

        async def mock_aget(experiment_id):
            nonlocal call_count
            call_count += 1
            if call_count <= 1:
                # First poll: WAITING_FOR_EXTERNAL
                return Experiment._from_dict(SAMPLE_EXTERNAL_AGENT_EXPERIMENT)
            else:
                # Second poll: COMPLETED (after agent ran)
                return Experiment._from_dict(
                    {**SAMPLE_EXTERNAL_AGENT_EXPERIMENT, "status": 4}
                )

        agent_calls = []

        async def my_agent(input_value):
            agent_calls.append(input_value)
            return "agent output"

        with patch.object(Experiment, "aget", side_effect=mock_aget):
            # First poll returns pending runs, second poll returns none (terminal)
            runs_call_count = 0

            async def mock_alist_runs(**kwargs):
                nonlocal runs_call_count
                runs_call_count += 1
                if runs_call_count == 1:
                    return (
                        [ExperimentRun.from_dict(SAMPLE_PENDING_RUN)],
                        MagicMock(total_pages=1),
                    )
                return ([], MagicMock(total_pages=1))

            with patch.object(exp, "alist_runs", side_effect=mock_alist_runs):
                # Mock record cache
                with patch.object(
                    exp,
                    "alist_records",
                    new_callable=AsyncMock,
                    return_value=(
                        [
                            MagicMock(
                                record_id="rec-001",
                                data={"text": "Hello world"},
                            )
                        ],
                        MagicMock(total_pages=1),
                    ),
                ):
                    # Mock aupdate_run
                    with patch.object(
                        exp,
                        "aupdate_run",
                        new_callable=AsyncMock,
                        return_value=ExperimentRun.from_dict(
                            {**SAMPLE_PENDING_RUN, "status": "completed"}
                        ),
                    ):
                        result = await exp.astart(
                            external_agents={"My Agent": my_agent},
                            trace_calls=False,
                        )

                        assert result.status == ExperimentStatus.COMPLETED
                        assert len(agent_calls) == 1
                        assert agent_calls[0] == "Hello world"

    @pytest.mark.asyncio
    @patch("lumenova_beacon.experiments.models.get_transport")
    @patch("lumenova_beacon.experiments.models.get_base_url")
    async def test_astart_starts_draft_then_polls(
        self, mock_base_url, mock_transport
    ):
        """DRAFT experiment is started via _post_start before polling."""
        mock_base_url.return_value = "https://api.example.com"
        mock_transport.return_value = _mock_transport()

        # _post_start returns QUEUED
        started_response = {**SAMPLE_EXTERNAL_AGENT_EXPERIMENT, "status": 2}
        mock_response = _mock_response(started_response)

        with patch("httpx.AsyncClient") as mock_client_cls:
            mock_client = AsyncMock()
            mock_client.post = AsyncMock(return_value=mock_response)
            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
            mock_client.__aexit__ = AsyncMock(return_value=False)
            mock_client_cls.return_value = mock_client

            exp = Experiment._from_dict(
                {**SAMPLE_EXTERNAL_AGENT_EXPERIMENT, "status": 1}  # DRAFT
            )

            # After _post_start, aget returns COMPLETED to exit polling
            completed_exp = Experiment._from_dict(
                {**SAMPLE_EXTERNAL_AGENT_EXPERIMENT, "status": 4}
            )
            with patch.object(
                Experiment, "aget", new_callable=AsyncMock
            ) as mock_aget:
                mock_aget.return_value = completed_exp

                result = await exp.astart(
                    external_agents={"My Agent": lambda x: x}
                )

                # Verify _post_start was called (POST to /start)
                mock_client.post.assert_called_once()
                assert result.status == ExperimentStatus.COMPLETED
