learningcircuit--local-deep-research/tests/database/test_thread_metrics_writer.py

"""Comprehensive tests for ThreadSafeMetricsWriter in database/thread_metrics.py."""

import threading
from unittest.mock import MagicMock, patch

import pytest

from local_deep_research.database.thread_metrics import (
    ThreadSafeMetricsWriter,
    metrics_writer,
)


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


@pytest.fixture
def writer():
    """Return a fresh ThreadSafeMetricsWriter instance."""
    return ThreadSafeMetricsWriter()


@pytest.fixture
def writer_with_password(writer):
    """Return a writer that already has a password stored for 'testuser'."""
    writer.set_user_password("testuser", "testpass")
    return writer


@pytest.fixture
def mock_db_manager():
    """Patch db_manager and yield (mock_db_manager, mock_session)."""
    with patch(
        "local_deep_research.database.thread_metrics.db_manager"
    ) as mock_db:
        mock_session = MagicMock()
        mock_db.create_thread_safe_session_for_metrics.return_value = (
            mock_session
        )
        yield mock_db, mock_session


# ===========================================================================
# set_user_password
# ===========================================================================


class TestSetUserPassword:
    """Tests for ThreadSafeMetricsWriter.set_user_password."""

    def test_stores_password_in_thread_local(self, writer):
        """set_user_password stores the password under _thread_local.passwords."""
        writer.set_user_password("alice", "secret123")

        assert writer._thread_local.passwords["alice"] == "secret123"

    def test_creates_passwords_dict_if_not_exists(self, writer):
        """Passwords dict is lazily created on first call."""
        assert not hasattr(writer._thread_local, "passwords")

        writer.set_user_password("alice", "secret")

        assert hasattr(writer._thread_local, "passwords")
        assert isinstance(writer._thread_local.passwords, dict)

    def test_can_store_multiple_users(self, writer):
        """Multiple users can be stored simultaneously."""
        writer.set_user_password("alice", "pass_a")
        writer.set_user_password("bob", "pass_b")
        writer.set_user_password("charlie", "pass_c")

        assert writer._thread_local.passwords == {
            "alice": "pass_a",
            "bob": "pass_b",
            "charlie": "pass_c",
        }

    def test_overwrites_existing_password(self, writer):
        """Setting a password for the same user replaces the old value."""
        writer.set_user_password("alice", "old_password")
        writer.set_user_password("alice", "new_password")

        assert writer._thread_local.passwords["alice"] == "new_password"
        assert len(writer._thread_local.passwords) == 1


# ===========================================================================
# get_session
# ===========================================================================


class TestGetSession:
    """Tests for ThreadSafeMetricsWriter.get_session."""

    # --- error paths -------------------------------------------------------

    def test_raises_when_no_passwords_dict(self, writer):
        """Raises ValueError when passwords dict has never been created."""
        with pytest.raises(ValueError, match="No password set"):
            with writer.get_session("testuser"):
                pass

    def test_raises_when_password_not_found_for_user(self, writer):
        """Raises ValueError when the requested user has no stored password."""
        writer.set_user_password("other_user", "other_pass")

        with pytest.raises(ValueError, match="No password available"):
            with writer.get_session("missing_user"):
                pass

    def test_raises_when_username_none_and_no_flask_context(self, writer):
        """Raises ValueError when username=None and Flask context is absent."""
        writer.set_user_password("alice", "pass")

        # The real code does `from flask import session as flask_session`
        # inside the function body. A RuntimeError simulates no app context.
        with pytest.raises((ValueError, RuntimeError)):
            with writer.get_session(username=None):
                pass

    def test_raises_when_db_manager_returns_none_session(
        self, writer_with_password
    ):
        """Raises ValueError when db_manager returns None for the session."""
        with patch(
            "local_deep_research.database.thread_metrics.db_manager"
        ) as mock_db:
            mock_db.create_thread_safe_session_for_metrics.return_value = None

            with pytest.raises(ValueError, match="Failed to create session"):
                with writer_with_password.get_session("testuser"):
                    pass

    # --- happy path --------------------------------------------------------

    def test_yields_session_on_success(
        self, writer_with_password, mock_db_manager
    ):
        """Successfully yields a session object when password is available."""
        _mock_db, mock_session = mock_db_manager

        with writer_with_password.get_session("testuser") as session:
            assert session is mock_session

    def test_creates_session_with_correct_args(
        self, writer_with_password, mock_db_manager
    ):
        """Passes correct username and password to db_manager."""
        mock_db, _mock_session = mock_db_manager

        with writer_with_password.get_session("testuser"):
            pass

        mock_db.create_thread_safe_session_for_metrics.assert_called_once_with(
            "testuser", "testpass"
        )

    def test_commits_session_on_success(
        self, writer_with_password, mock_db_manager
    ):
        """Session is committed when context manager exits normally."""
        _mock_db, mock_session = mock_db_manager

        with writer_with_password.get_session("testuser"):
            pass

        mock_session.commit.assert_called_once()
        mock_session.rollback.assert_not_called()

    def test_rollback_on_exception(self, writer_with_password, mock_db_manager):
        """Session is rolled back when an exception occurs inside the block."""
        _mock_db, mock_session = mock_db_manager

        with pytest.raises(RuntimeError, match="boom"):
            with writer_with_password.get_session("testuser"):
                raise RuntimeError("boom")

        mock_session.rollback.assert_called_once()
        mock_session.commit.assert_not_called()

    def test_closes_session_on_success(
        self, writer_with_password, mock_db_manager
    ):
        """Session is closed in the finally block after normal exit."""
        _mock_db, mock_session = mock_db_manager

        with writer_with_password.get_session("testuser"):
            pass

        mock_session.close.assert_called_once()

    def test_closes_session_on_exception(
        self, writer_with_password, mock_db_manager
    ):
        """Session is closed in the finally block even after an exception."""
        _mock_db, mock_session = mock_db_manager

        with pytest.raises(RuntimeError):
            with writer_with_password.get_session("testuser"):
                raise RuntimeError("failure")

        mock_session.close.assert_called_once()

    def test_closes_session_when_db_manager_returns_none(
        self, writer_with_password
    ):
        """When db_manager returns None, session is not closed (it is None)."""
        with patch(
            "local_deep_research.database.thread_metrics.db_manager"
        ) as mock_db:
            mock_db.create_thread_safe_session_for_metrics.return_value = None

            with pytest.raises(ValueError, match="Failed to create session"):
                with writer_with_password.get_session("testuser"):
                    pass

            # session was None, so close should never be called on it


# ===========================================================================
# write_token_metrics
# ===========================================================================


class TestWriteTokenMetrics:
    """Tests for ThreadSafeMetricsWriter.write_token_metrics."""

    def _run_write(self, writer, mock_db_manager, token_data, research_id=42):
        """Helper: call write_token_metrics with mocked session and TokenUsage."""
        mock_db, mock_session = mock_db_manager
        mock_token_cls = MagicMock()

        with patch(
            "local_deep_research.database.models.TokenUsage",
            mock_token_cls,
        ):
            writer.write_token_metrics("testuser", research_id, token_data)

        return mock_token_cls, mock_session

    # --- field mapping -----------------------------------------------------

    def test_creates_token_usage_with_correct_fields(
        self, writer_with_password, mock_db_manager
    ):
        """TokenUsage is constructed with all expected fields from token_data."""
        token_data = {
            "model_name": "gpt-4",
            "provider": "openai",
            "prompt_tokens": 100,
            "completion_tokens": 50,
            "research_query": "test query",
            "research_mode": "deep",
            "research_phase": "search",
            "search_iteration": 2,
            "response_time_ms": 1234,
            "success_status": "success",
            "error_type": None,
            "search_engines_planned": ["google", "bing"],
            "search_engine_selected": "google",
            "calling_file": "main.py",
            "calling_function": "run_search",
            "call_stack": ["a", "b"],
            "context_limit": 4096,
            "context_truncated": True,
            "tokens_truncated": 500,
            "truncation_ratio": 0.12,
            "ollama_prompt_eval_count": 90,
            "ollama_eval_count": 45,
            "ollama_total_duration": 999999,
            "ollama_load_duration": 111111,
            "ollama_prompt_eval_duration": 222222,
            "ollama_eval_duration": 333333,
        }

        mock_token_cls, _ = self._run_write(
            writer_with_password, mock_db_manager, token_data, research_id=7
        )

        call_kwargs = mock_token_cls.call_args[1]

        assert call_kwargs["research_id"] == 7
        assert call_kwargs["model_name"] == "gpt-4"
        assert call_kwargs["model_provider"] == "openai"
        assert call_kwargs["prompt_tokens"] == 100
        assert call_kwargs["completion_tokens"] == 50
        assert call_kwargs["total_tokens"] == 150
        assert call_kwargs["research_query"] == "test query"
        assert call_kwargs["research_mode"] == "deep"
        assert call_kwargs["research_phase"] == "search"
        assert call_kwargs["search_iteration"] == 2
        assert call_kwargs["response_time_ms"] == 1234
        assert call_kwargs["success_status"] == "success"
        assert call_kwargs["error_type"] is None
        assert call_kwargs["search_engines_planned"] == ["google", "bing"]
        assert call_kwargs["search_engine_selected"] == "google"
        assert call_kwargs["calling_file"] == "main.py"
        assert call_kwargs["calling_function"] == "run_search"
        assert call_kwargs["call_stack"] == ["a", "b"]
        assert call_kwargs["context_limit"] == 4096
        assert call_kwargs["context_truncated"] is True
        assert call_kwargs["tokens_truncated"] == 500
        assert call_kwargs["truncation_ratio"] == 0.12
        assert call_kwargs["ollama_prompt_eval_count"] == 90
        assert call_kwargs["ollama_eval_count"] == 45
        assert call_kwargs["ollama_total_duration"] == 999999
        assert call_kwargs["ollama_load_duration"] == 111111
        assert call_kwargs["ollama_prompt_eval_duration"] == 222222
        assert call_kwargs["ollama_eval_duration"] == 333333

    def test_adds_record_to_session(
        self, writer_with_password, mock_db_manager
    ):
        """The TokenUsage record is added to the session via session.add()."""
        token_data = {
            "model_name": "gpt-4",
            "provider": "openai",
            "prompt_tokens": 10,
            "completion_tokens": 5,
        }

        mock_token_cls, mock_session = self._run_write(
            writer_with_password, mock_db_manager, token_data
        )

        mock_session.add.assert_called_once_with(mock_token_cls.return_value)

    def test_total_tokens_calculated_as_sum(
        self, writer_with_password, mock_db_manager
    ):
        """total_tokens equals prompt_tokens + completion_tokens."""
        token_data = {
            "prompt_tokens": 200,
            "completion_tokens": 75,
        }

        mock_token_cls, _ = self._run_write(
            writer_with_password, mock_db_manager, token_data
        )

        call_kwargs = mock_token_cls.call_args[1]
        assert call_kwargs["total_tokens"] == 275
        assert (
            call_kwargs["total_tokens"]
            == call_kwargs["prompt_tokens"] + call_kwargs["completion_tokens"]
        )

    # --- defaults ----------------------------------------------------------

    def test_default_prompt_tokens_zero(
        self, writer_with_password, mock_db_manager
    ):
        """prompt_tokens defaults to 0 when not provided."""
        token_data = {"completion_tokens": 30}

        mock_token_cls, _ = self._run_write(
            writer_with_password, mock_db_manager, token_data
        )

        call_kwargs = mock_token_cls.call_args[1]
        assert call_kwargs["prompt_tokens"] == 0

    def test_default_completion_tokens_zero(
        self, writer_with_password, mock_db_manager
    ):
        """completion_tokens defaults to 0 when not provided."""
        token_data = {"prompt_tokens": 50}

        mock_token_cls, _ = self._run_write(
            writer_with_password, mock_db_manager, token_data
        )

        call_kwargs = mock_token_cls.call_args[1]
        assert call_kwargs["completion_tokens"] == 0

    def test_default_success_status(
        self, writer_with_password, mock_db_manager
    ):
        """success_status defaults to 'success' when not provided."""
        token_data = {}

        mock_token_cls, _ = self._run_write(
            writer_with_password, mock_db_manager, token_data
        )

        call_kwargs = mock_token_cls.call_args[1]
        assert call_kwargs["success_status"] == "success"

    def test_default_context_truncated_false(
        self, writer_with_password, mock_db_manager
    ):
        """context_truncated defaults to False when not provided."""
        token_data = {}

        mock_token_cls, _ = self._run_write(
            writer_with_password, mock_db_manager, token_data
        )

        call_kwargs = mock_token_cls.call_args[1]
        assert call_kwargs["context_truncated"] is False

    def test_default_total_tokens_zero_when_both_missing(
        self, writer_with_password, mock_db_manager
    ):
        """total_tokens is 0 when both prompt_tokens and completion_tokens are absent."""
        token_data = {}

        mock_token_cls, _ = self._run_write(
            writer_with_password, mock_db_manager, token_data
        )

        call_kwargs = mock_token_cls.call_args[1]
        assert call_kwargs["total_tokens"] == 0

    # --- missing optional fields -------------------------------------------

    def test_missing_optional_fields_are_none(
        self, writer_with_password, mock_db_manager
    ):
        """Optional fields not present in token_data resolve to None via .get()."""
        token_data = {}

        mock_token_cls, _ = self._run_write(
            writer_with_password, mock_db_manager, token_data
        )

        call_kwargs = mock_token_cls.call_args[1]
        assert call_kwargs["model_name"] is None
        assert call_kwargs["model_provider"] is None
        assert call_kwargs["research_query"] is None
        assert call_kwargs["research_mode"] is None
        assert call_kwargs["research_phase"] is None
        assert call_kwargs["search_iteration"] is None
        assert call_kwargs["response_time_ms"] is None
        assert call_kwargs["error_type"] is None
        assert call_kwargs["search_engines_planned"] is None
        assert call_kwargs["search_engine_selected"] is None
        assert call_kwargs["calling_file"] is None
        assert call_kwargs["calling_function"] is None
        assert call_kwargs["call_stack"] is None
        assert call_kwargs["context_limit"] is None
        assert call_kwargs["tokens_truncated"] is None
        assert call_kwargs["truncation_ratio"] is None
        assert call_kwargs["ollama_prompt_eval_count"] is None
        assert call_kwargs["ollama_eval_count"] is None
        assert call_kwargs["ollama_total_duration"] is None
        assert call_kwargs["ollama_load_duration"] is None
        assert call_kwargs["ollama_prompt_eval_duration"] is None
        assert call_kwargs["ollama_eval_duration"] is None

    def test_research_id_passed_directly(
        self, writer_with_password, mock_db_manager
    ):
        """research_id comes from the function argument, not from token_data."""
        token_data = {"model_name": "test-model"}

        mock_token_cls, _ = self._run_write(
            writer_with_password, mock_db_manager, token_data, research_id=999
        )

        call_kwargs = mock_token_cls.call_args[1]
        assert call_kwargs["research_id"] == 999

    def test_research_id_none(self, writer_with_password, mock_db_manager):
        """research_id can be None."""
        token_data = {"model_name": "test-model"}

        mock_token_cls, _ = self._run_write(
            writer_with_password, mock_db_manager, token_data, research_id=None
        )

        call_kwargs = mock_token_cls.call_args[1]
        assert call_kwargs["research_id"] is None

    # --- ModelUsage upsert -------------------------------------------------

    def test_creates_model_usage_when_none_exists(
        self, writer_with_password, mock_db_manager
    ):
        """New ModelUsage record is created when none exists."""
        token_data = {
            "model_name": "gpt-4",
            "provider": "openai",
            "prompt_tokens": 100,
            "completion_tokens": 50,
        }

        mock_db, mock_session = mock_db_manager
        mock_session.query.return_value.filter_by.return_value.first.return_value = None

        with patch(
            "local_deep_research.database.models.TokenUsage", MagicMock()
        ):
            with patch(
                "local_deep_research.database.models.ModelUsage"
            ) as mock_model_cls:
                writer_with_password.write_token_metrics(
                    "testuser", 1, token_data
                )

                mock_model_cls.assert_called_once_with(
                    model_name="gpt-4",
                    model_provider="openai",
                    total_tokens=150,
                    total_calls=1,
                )
                assert mock_session.add.call_count == 2

    def test_updates_existing_model_usage(
        self, writer_with_password, mock_db_manager
    ):
        """Existing ModelUsage record is incremented."""
        token_data = {
            "model_name": "gpt-4",
            "provider": "openai",
            "prompt_tokens": 100,
            "completion_tokens": 50,
        }

        mock_db, mock_session = mock_db_manager
        existing = MagicMock()
        existing.total_tokens = 1000
        existing.total_calls = 5
        mock_session.query.return_value.filter_by.return_value.first.return_value = existing

        with patch(
            "local_deep_research.database.models.TokenUsage", MagicMock()
        ):
            writer_with_password.write_token_metrics("testuser", 1, token_data)

            assert existing.total_tokens == 1150
            assert existing.total_calls == 6


# ===========================================================================
# Global instance
# ===========================================================================


class TestGlobalInstance:
    """Tests for the module-level metrics_writer singleton."""

    def test_metrics_writer_exists(self):
        """A global metrics_writer instance is importable."""
        assert metrics_writer is not None

    def test_metrics_writer_is_correct_type(self):
        """The global instance is a ThreadSafeMetricsWriter."""
        assert isinstance(metrics_writer, ThreadSafeMetricsWriter)

    def test_metrics_writer_has_thread_local(self):
        """The global instance has _thread_local attribute."""
        assert hasattr(metrics_writer, "_thread_local")
        assert isinstance(metrics_writer._thread_local, threading.local)


# ===========================================================================
# Thread isolation
# ===========================================================================


class TestThreadIsolation:
    """Verify that thread-local storage is genuinely isolated between threads."""

    def test_passwords_not_shared_across_threads(self, writer):
        """Password set in one thread is invisible in another thread."""
        results = {}

        def thread_a():
            writer.set_user_password("alice", "alice_pass")
            results["a_passwords"] = dict(writer._thread_local.passwords)

        def thread_b():
            # Wait for thread_a to finish writing
            import time

            time.sleep(0.05)
            results["b_has_passwords"] = hasattr(
                writer._thread_local, "passwords"
            )

        t_a = threading.Thread(target=thread_a)
        t_b = threading.Thread(target=thread_b)
        t_a.start()
        t_b.start()
        t_a.join()
        t_b.join()

        assert results["a_passwords"] == {"alice": "alice_pass"}
        assert results["b_has_passwords"] is False

    def test_each_thread_has_own_password_store(self, writer):
        """Two threads can independently store passwords without interference."""
        results = {}
        barrier = threading.Barrier(2)

        def thread_one():
            writer.set_user_password("user_one", "pass_one")
            barrier.wait()
            results["one"] = dict(writer._thread_local.passwords)

        def thread_two():
            writer.set_user_password("user_two", "pass_two")
            barrier.wait()
            results["two"] = dict(writer._thread_local.passwords)

        t1 = threading.Thread(target=thread_one)
        t2 = threading.Thread(target=thread_two)
        t1.start()
        t2.start()
        t1.join()
        t2.join()

        assert results["one"] == {"user_one": "pass_one"}
        assert results["two"] == {"user_two": "pass_two"}
        # Neither thread sees the other's password
        assert "user_two" not in results["one"]
        assert "user_one" not in results["two"]

    def test_get_session_fails_in_different_thread(self, writer):
        """get_session in a thread that hasn't called set_user_password raises."""
        writer.set_user_password("alice", "pass")

        errors = []

        def other_thread():
            try:
                with writer.get_session("alice"):
                    pass
            except ValueError as e:
                errors.append(str(e))

        t = threading.Thread(target=other_thread)
        t.start()
        t.join()

        assert len(errors) == 1
        assert "No password set" in errors[0]