Shepherd MCP

"""Tests for provider base classes and AIOBS provider."""

from unittest.mock import Mock, patch

import pytest

from shepherd_mcp.models.aiobs import (
    Event,
    FunctionEvent,
    Session,
    SessionsResponse,
)
from shepherd_mcp.providers.aiobs import (
    AIOBSClient,
    eval_is_failed,
    filter_sessions,
    parse_date,
    session_has_errors,
    session_has_failed_evals,
    session_has_function,
    session_has_model,
    session_has_provider,
    session_matches_labels,
    session_matches_query,
)
from shepherd_mcp.providers.base import (
    AuthenticationError,
    NotFoundError,
    ProviderError,
    RateLimitError,
)

# ============================================================================
# Base Provider Tests
# ============================================================================


class TestExceptions:
    """Tests for provider exceptions."""

    def test_provider_error_is_exception(self):
        with pytest.raises(ProviderError):
            raise ProviderError("test error")

    def test_authentication_error_is_provider_error(self):
        with pytest.raises(ProviderError):
            raise AuthenticationError("auth failed")

    def test_not_found_error_is_provider_error(self):
        with pytest.raises(ProviderError):
            raise NotFoundError("not found")

    def test_rate_limit_error_is_provider_error(self):
        with pytest.raises(ProviderError):
            raise RateLimitError("rate limited")


# ============================================================================
# AIOBS Provider Tests
# ============================================================================


class TestAIOBSClientInit:
    """Tests for AIOBSClient initialization."""

    def test_missing_api_key_raises_error(self):
        with patch.dict("os.environ", {}, clear=True):
            with pytest.raises(AuthenticationError) as exc_info:
                AIOBSClient()
            assert "AIOBS_API_KEY" in str(exc_info.value)

    def test_init_with_env_var(self):
        with patch.dict("os.environ", {"AIOBS_API_KEY": "test-key"}):
            client = AIOBSClient()
            assert client.api_key == "test-key"
            assert "shepherd-api" in client.endpoint
            client.close()

    def test_init_with_explicit_key(self):
        client = AIOBSClient(api_key="explicit-key")
        assert client.api_key == "explicit-key"
        client.close()

    def test_init_with_custom_endpoint(self):
        client = AIOBSClient(
            api_key="test-key",
            endpoint="https://custom-api.example.com/",
        )
        assert client.endpoint == "https://custom-api.example.com"  # trailing slash removed
        client.close()

    def test_provider_name(self):
        client = AIOBSClient(api_key="test-key")
        assert client.name == "aiobs"
        client.close()


class TestAIOBSClientContextManager:
    """Tests for AIOBSClient context manager."""

    def test_context_manager(self):
        with AIOBSClient(api_key="test-key") as client:
            assert client.name == "aiobs"


class TestAIOBSClientErrorHandling:
    """Tests for AIOBSClient error handling."""

    def setup_method(self):
        self.client = AIOBSClient(api_key="test-key")

    def teardown_method(self):
        self.client.close()

    def test_401_raises_authentication_error(self):
        mock_response = Mock()
        mock_response.status_code = 401
        mock_response.json.return_value = {"detail": "Invalid API key"}
        with pytest.raises(AuthenticationError) as exc_info:
            self.client._handle_error_response(mock_response)
        assert "Invalid API key" in str(exc_info.value)

    def test_404_raises_not_found_error(self):
        mock_response = Mock()
        mock_response.status_code = 404
        mock_response.json.return_value = {"detail": "Session not found"}
        with pytest.raises(NotFoundError) as exc_info:
            self.client._handle_error_response(mock_response)
        assert "Session not found" in str(exc_info.value)

    def test_500_raises_provider_error(self):
        mock_response = Mock()
        mock_response.status_code = 500
        mock_response.json.return_value = {"detail": "Internal error"}
        with pytest.raises(ProviderError) as exc_info:
            self.client._handle_error_response(mock_response)
        assert "Internal error" in str(exc_info.value)


# ============================================================================
# Parse Date Tests
# ============================================================================


class TestParseDate:
    """Tests for parse_date function."""

    def test_date_only(self):
        result = parse_date("2025-01-15")
        # Should be a Unix timestamp
        assert isinstance(result, float)
        assert result > 0

    def test_datetime_with_time(self):
        result = parse_date("2025-01-15 12:30:45")
        assert isinstance(result, float)

    def test_datetime_iso_format(self):
        result = parse_date("2025-01-15T12:30:45")
        assert isinstance(result, float)

    def test_datetime_short(self):
        result = parse_date("2025-01-15 12:30")
        assert isinstance(result, float)

    def test_invalid_format_raises_error(self):
        with pytest.raises(ValueError) as exc_info:
            parse_date("invalid-date")
        assert "Invalid date format" in str(exc_info.value)


# ============================================================================
# Session Filtering Tests
# ============================================================================


def make_session(
    id: str = "test-session",
    name: str = "Test Session",
    started_at: float = 1735689600.0,
    labels: dict = None,
    meta: dict = None,
) -> Session:
    """Helper to create Session objects for testing."""
    return Session(
        id=id,
        name=name,
        started_at=started_at,
        ended_at=started_at + 60,
        labels=labels or {},
        meta=meta or {},
    )


def make_event(
    session_id: str = "test-session",
    provider: str = "openai",
    model: str = "gpt-4o-mini",
    error: str = None,
    evaluations: list = None,
) -> Event:
    """Helper to create Event objects for testing."""
    return Event(
        provider=provider,
        api="chat.completions.create",
        request={"model": model},
        response=None,
        error=error,
        started_at=1735689600.0,
        ended_at=1735689601.0,
        duration_ms=1000.0,
        span_id="span-1",
        session_id=session_id,
        evaluations=evaluations or [],
    )


def make_function_event(
    session_id: str = "test-session",
    name: str = "my_function",
    module: str = "my_module",
    error: str = None,
    evaluations: list = None,
) -> FunctionEvent:
    """Helper to create FunctionEvent objects for testing."""
    return FunctionEvent(
        provider="function",
        api="call",
        name=name,
        module=module,
        error=error,
        started_at=1735689600.0,
        ended_at=1735689601.0,
        duration_ms=1000.0,
        span_id="span-1",
        session_id=session_id,
        evaluations=evaluations or [],
    )


class TestSessionMatchesQuery:
    """Tests for session_matches_query."""

    def test_matches_id(self):
        session = make_session(id="abc-123-def")
        assert session_matches_query(session, "abc-123") is True

    def test_matches_name(self):
        session = make_session(name="Production Run")
        assert session_matches_query(session, "production") is True

    def test_matches_label_value(self):
        session = make_session(labels={"env": "staging"})
        assert session_matches_query(session, "staging") is True

    def test_matches_meta_value(self):
        session = make_session(meta={"user": "john@example.com"})
        assert session_matches_query(session, "john") is True

    def test_no_match(self):
        session = make_session(id="abc", name="Test", labels={}, meta={})
        assert session_matches_query(session, "xyz") is False


class TestSessionMatchesLabels:
    """Tests for session_matches_labels."""

    def test_matches_all_labels(self):
        session = make_session(labels={"env": "prod", "team": "ml"})
        assert session_matches_labels(session, {"env": "prod", "team": "ml"}) is True

    def test_matches_subset(self):
        session = make_session(labels={"env": "prod", "team": "ml"})
        assert session_matches_labels(session, {"env": "prod"}) is True

    def test_missing_label(self):
        session = make_session(labels={"env": "prod"})
        assert session_matches_labels(session, {"team": "ml"}) is False

    def test_wrong_value(self):
        session = make_session(labels={"env": "prod"})
        assert session_matches_labels(session, {"env": "staging"}) is False


class TestSessionHasProvider:
    """Tests for session_has_provider."""

    def test_matches_provider(self):
        session = make_session(id="s1")
        events = [make_event(session_id="s1", provider="openai")]
        assert session_has_provider(session, events, [], "openai") is True

    def test_case_insensitive(self):
        session = make_session(id="s1")
        events = [make_event(session_id="s1", provider="OpenAI")]
        assert session_has_provider(session, events, [], "OPENAI") is True

    def test_no_match(self):
        session = make_session(id="s1")
        events = [make_event(session_id="s1", provider="anthropic")]
        assert session_has_provider(session, events, [], "openai") is False

    def test_matches_in_function_events(self):
        session = make_session(id="s1")
        fn_events = [make_function_event(session_id="s1")]
        fn_events[0].provider = "custom"
        assert session_has_provider(session, [], fn_events, "custom") is True


class TestSessionHasModel:
    """Tests for session_has_model."""

    def test_matches_model(self):
        session = make_session(id="s1")
        events = [make_event(session_id="s1", model="gpt-4o-mini")]
        assert session_has_model(session, events, "gpt-4o") is True

    def test_case_insensitive(self):
        session = make_session(id="s1")
        events = [make_event(session_id="s1", model="GPT-4o-Mini")]
        assert session_has_model(session, events, "gpt-4o-mini") is True

    def test_no_match(self):
        session = make_session(id="s1")
        events = [make_event(session_id="s1", model="claude-3")]
        assert session_has_model(session, events, "gpt-4") is False


class TestSessionHasErrors:
    """Tests for session_has_errors."""

    def test_event_with_error(self):
        session = make_session(id="s1")
        events = [make_event(session_id="s1", error="API error")]
        assert session_has_errors(session, events, []) is True

    def test_function_event_with_error(self):
        session = make_session(id="s1")
        fn_events = [make_function_event(session_id="s1", error="Function failed")]
        assert session_has_errors(session, [], fn_events) is True

    def test_no_errors(self):
        session = make_session(id="s1")
        events = [make_event(session_id="s1")]
        fn_events = [make_function_event(session_id="s1")]
        assert session_has_errors(session, events, fn_events) is False


class TestSessionHasFunction:
    """Tests for session_has_function."""

    def test_matches_function_name(self):
        session = make_session(id="s1")
        fn_events = [make_function_event(session_id="s1", name="process_data")]
        assert session_has_function(session, fn_events, "process") is True

    def test_matches_module_name(self):
        session = make_session(id="s1")
        fn_events = [make_function_event(session_id="s1", module="data_pipeline")]
        assert session_has_function(session, fn_events, "pipeline") is True

    def test_no_match(self):
        session = make_session(id="s1")
        fn_events = [make_function_event(session_id="s1", name="func_a")]
        assert session_has_function(session, fn_events, "func_b") is False


class TestEvalIsFailed:
    """Tests for eval_is_failed."""

    def test_passed_false(self):
        assert eval_is_failed({"passed": False}) is True

    def test_passed_true(self):
        assert eval_is_failed({"passed": True}) is False

    def test_result_false(self):
        assert eval_is_failed({"result": False}) is True

    def test_status_failed(self):
        assert eval_is_failed({"status": "failed"}) is True
        assert eval_is_failed({"status": "fail"}) is True
        assert eval_is_failed({"status": "error"}) is True

    def test_success_false(self):
        assert eval_is_failed({"success": False}) is True

    def test_not_dict(self):
        assert eval_is_failed("not a dict") is False
        assert eval_is_failed(None) is False


class TestSessionHasFailedEvals:
    """Tests for session_has_failed_evals."""

    def test_event_with_failed_eval(self):
        session = make_session(id="s1")
        events = [make_event(session_id="s1", evaluations=[{"passed": False}])]
        assert session_has_failed_evals(session, events, []) is True

    def test_function_event_with_failed_eval(self):
        session = make_session(id="s1")
        fn_events = [make_function_event(session_id="s1", evaluations=[{"passed": False}])]
        assert session_has_failed_evals(session, [], fn_events) is True

    def test_no_failed_evals(self):
        session = make_session(id="s1")
        events = [make_event(session_id="s1", evaluations=[{"passed": True}])]
        assert session_has_failed_evals(session, events, []) is False


class TestFilterSessions:
    """Tests for filter_sessions function."""

    def setup_method(self):
        self.sessions = [
            make_session(id="s1", name="Production", labels={"env": "prod"}),
            make_session(id="s2", name="Staging", labels={"env": "staging"}),
            make_session(id="s3", name="Development", labels={"env": "dev"}),
        ]
        self.events = [
            make_event(session_id="s1", provider="openai", model="gpt-4"),
            make_event(session_id="s2", provider="anthropic", model="claude-3"),
            make_event(session_id="s3", provider="openai", model="gpt-3.5"),
        ]
        self.response = SessionsResponse(
            sessions=self.sessions,
            events=self.events,
            function_events=[],
        )

    def test_no_filters(self):
        result = filter_sessions(self.response)
        assert len(result.sessions) == 3

    def test_filter_by_query(self):
        result = filter_sessions(self.response, query="prod")
        assert len(result.sessions) == 1
        assert result.sessions[0].id == "s1"

    def test_filter_by_labels(self):
        result = filter_sessions(self.response, labels={"env": "staging"})
        assert len(result.sessions) == 1
        assert result.sessions[0].id == "s2"

    def test_filter_by_provider(self):
        result = filter_sessions(self.response, provider="openai")
        assert len(result.sessions) == 2
        assert "s1" in [s.id for s in result.sessions]
        assert "s3" in [s.id for s in result.sessions]

    def test_filter_by_model(self):
        result = filter_sessions(self.response, model="gpt-4")
        assert len(result.sessions) == 1
        assert result.sessions[0].id == "s1"

    def test_events_filtered_with_sessions(self):
        result = filter_sessions(self.response, query="prod")
        # Events should only include those from matching sessions
        assert len(result.events) == 1
        assert result.events[0].session_id == "s1"