"""Tests for test_runner.validate.baseline_loader (StageBaselineLoader)."""

from __future__ import annotations

import json
from unittest.mock import MagicMock, patch

import pytest

from test_runner.validate.baseline_loader import StageBaselineLoader


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _make_stage_row(
    procedure_name: str = "dbo.GetProducts",
    params: dict | None = None,
    captured_at: str = "2026-01-01T00:00:00Z",
    success: bool = True,
    error: str | None = None,
    row_counts: list | str | None = None,
) -> dict:
    row: dict = {
        "PROCEDURE_NAME": procedure_name,
        "PARAMS": json.dumps(params or {}),
        "CAPTURED_AT": captured_at,
        "SUCCESS": success,
        "ERROR": error,
    }
    if row_counts is not None:
        row["ROW_COUNTS"] = json.dumps(row_counts) if isinstance(row_counts, list) else row_counts
    return row


def _make_loader(stage_rows: list[dict]) -> StageBaselineLoader:
    connector = MagicMock()
    connector.execute_query.return_value = stage_rows
    connector.execute_statement = MagicMock()
    loader = StageBaselineLoader(connector, "MYDB")
    return loader


# ---------------------------------------------------------------------------
# TestStageBaselineLoaderRowCounts
# ---------------------------------------------------------------------------

class TestStageBaselineLoaderRowCounts:
    def test_row_counts_populated_from_stage(self):
        rows = [_make_stage_row(row_counts=[15])]
        loader = _make_loader(rows)
        baselines = loader.load()
        assert len(baselines) == 1
        assert baselines[0].row_counts == [15]

    def test_row_counts_defaults_to_empty_when_absent(self):
        row = _make_stage_row()
        # ROW_COUNTS key not present
        assert "ROW_COUNTS" not in row
        loader = _make_loader([row])
        baselines = loader.load()
        assert baselines[0].row_counts == []

    def test_multiple_result_sets_row_counts(self):
        rows = [_make_stage_row(row_counts=[10, 5])]
        loader = _make_loader(rows)
        baselines = loader.load()
        assert baselines[0].row_counts == [10, 5]

    def test_row_counts_null_value_defaults_to_empty(self):
        row = _make_stage_row()
        row["ROW_COUNTS"] = None
        loader = _make_loader([row])
        baselines = loader.load()
        assert baselines[0].row_counts == []

    def test_procedure_name_still_populated(self):
        rows = [_make_stage_row(procedure_name="Reports.GetTopProducts", row_counts=[20])]
        loader = _make_loader(rows)
        baselines = loader.load()
        assert baselines[0].procedure == "Reports.GetTopProducts"

    def test_stage_query_includes_row_counts_column(self):
        loader = _make_loader([])
        loader.load()
        call_args = loader._connector.execute_query.call_args
        query = call_args[0][0]
        assert "row_counts" in query.lower()
