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

"""Tests for database/thread_metrics.py."""

import pytest
import threading
from unittest.mock import Mock, patch


class TestThreadSafeMetricsWriter:
    """Tests for ThreadSafeMetricsWriter class."""

    def test_init_creates_thread_local_storage(self):
        """Test that initialization creates thread-local storage."""
        from local_deep_research.database.thread_metrics import (
            ThreadSafeMetricsWriter,
        )

        writer = ThreadSafeMetricsWriter()
        assert hasattr(writer, "_thread_local")
        assert isinstance(writer._thread_local, threading.local)

    def test_set_user_password_stores_in_thread_local(self):
        """Test that set_user_password stores password in thread-local storage."""
        from local_deep_research.database.thread_metrics import (
            ThreadSafeMetricsWriter,
        )

        writer = ThreadSafeMetricsWriter()
        writer.set_user_password("testuser", "testpass")

        assert hasattr(writer._thread_local, "passwords")
        assert writer._thread_local.passwords["testuser"] == "testpass"

    def test_set_user_password_creates_dict_if_missing(self):
        """Test that set_user_password creates passwords dict if it doesn't exist."""
        from local_deep_research.database.thread_metrics import (
            ThreadSafeMetricsWriter,
        )

        writer = ThreadSafeMetricsWriter()
        # Initially no passwords dict
        assert not hasattr(writer._thread_local, "passwords")

        writer.set_user_password("user1", "pass1")

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

    def test_set_user_password_supports_multiple_users(self):
        """Test that multiple users can have passwords stored."""
        from local_deep_research.database.thread_metrics import (
            ThreadSafeMetricsWriter,
        )

        writer = ThreadSafeMetricsWriter()
        writer.set_user_password("user1", "pass1")
        writer.set_user_password("user2", "pass2")

        assert writer._thread_local.passwords["user1"] == "pass1"
        assert writer._thread_local.passwords["user2"] == "pass2"

    def test_get_session_raises_when_no_password_set(self):
        """Test that get_session raises error when no password is set."""
        from local_deep_research.database.thread_metrics import (
            ThreadSafeMetricsWriter,
        )

        writer = ThreadSafeMetricsWriter()

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

    def test_get_session_raises_when_user_password_missing(self):
        """Test that get_session raises error when user's password is missing."""
        from local_deep_research.database.thread_metrics import (
            ThreadSafeMetricsWriter,
        )

        writer = ThreadSafeMetricsWriter()
        writer.set_user_password("otheruser", "otherpass")

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

    def test_get_session_creates_session_with_password(self):
        """Test that get_session creates session with stored password."""
        from local_deep_research.database.thread_metrics import (
            ThreadSafeMetricsWriter,
        )

        writer = ThreadSafeMetricsWriter()
        writer.set_user_password("testuser", "testpass")

        with patch(
            "local_deep_research.database.thread_metrics.db_manager"
        ) as mock_db:
            mock_session = Mock()
            mock_db.create_thread_safe_session_for_metrics.return_value = (
                mock_session
            )

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

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

    def test_get_session_commits_on_success(self):
        """Test that get_session commits the session on successful exit."""
        from local_deep_research.database.thread_metrics import (
            ThreadSafeMetricsWriter,
        )

        writer = ThreadSafeMetricsWriter()
        writer.set_user_password("testuser", "testpass")

        with patch(
            "local_deep_research.database.thread_metrics.db_manager"
        ) as mock_db:
            mock_session = Mock()
            mock_db.create_thread_safe_session_for_metrics.return_value = (
                mock_session
            )

            with writer.get_session("testuser"):
                pass

            mock_session.commit.assert_called_once()
            mock_session.close.assert_called_once()

    def test_get_session_rollbacks_on_error(self):
        """Test that get_session rolls back on error."""
        from local_deep_research.database.thread_metrics import (
            ThreadSafeMetricsWriter,
        )

        writer = ThreadSafeMetricsWriter()
        writer.set_user_password("testuser", "testpass")

        with patch(
            "local_deep_research.database.thread_metrics.db_manager"
        ) as mock_db:
            mock_session = Mock()
            mock_db.create_thread_safe_session_for_metrics.return_value = (
                mock_session
            )

            with pytest.raises(RuntimeError):
                with writer.get_session("testuser"):
                    raise RuntimeError("Test error")

            mock_session.rollback.assert_called_once()
            mock_session.close.assert_called_once()

    def test_get_session_raises_when_session_creation_fails(self):
        """Test that get_session raises error when session creation fails."""
        from local_deep_research.database.thread_metrics import (
            ThreadSafeMetricsWriter,
        )

        writer = ThreadSafeMetricsWriter()
        writer.set_user_password("testuser", "testpass")

        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.get_session("testuser"):
                    pass


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

    def test_write_token_metrics_creates_token_usage_record(self):
        """Test that write_token_metrics creates a TokenUsage record."""
        from local_deep_research.database.thread_metrics import (
            ThreadSafeMetricsWriter,
        )

        writer = ThreadSafeMetricsWriter()
        writer.set_user_password("testuser", "testpass")

        mock_session = Mock()
        mock_token_usage_class = Mock()

        with patch(
            "local_deep_research.database.thread_metrics.db_manager"
        ) as mock_db:
            mock_db.create_thread_safe_session_for_metrics.return_value = (
                mock_session
            )

            with patch(
                "local_deep_research.database.models.TokenUsage",
                mock_token_usage_class,
            ):
                with patch("local_deep_research.database.models.ModelUsage"):
                    mock_session.query.return_value.filter_by.return_value.first.return_value = None

                    token_data = {
                        "model_name": "gpt-4",
                        "provider": "openai",
                        "prompt_tokens": 100,
                        "completion_tokens": 50,
                    }

                    writer.write_token_metrics("testuser", 123, token_data)

                    # Verify TokenUsage was created and added to session
                    mock_token_usage_class.assert_called_once()
                    assert mock_session.add.call_count == 2


class TestGlobalInstance:
    """Tests for the global metrics_writer instance."""

    def test_global_instance_is_correct_type(self):
        """Test that global instance is ThreadSafeMetricsWriter."""
        from local_deep_research.database.thread_metrics import (
            metrics_writer,
            ThreadSafeMetricsWriter,
        )

        assert isinstance(metrics_writer, ThreadSafeMetricsWriter)


class TestClearPasswords:
    """Tests for clear_passwords() and its integration with cleanup_current_thread()."""

    def test_clear_passwords_empties_dict(self):
        from local_deep_research.database.thread_metrics import (
            ThreadSafeMetricsWriter,
        )

        writer = ThreadSafeMetricsWriter()
        writer.set_user_password("userA", "passA")
        writer.set_user_password("userB", "passB")
        assert writer._thread_local.passwords == {
            "userA": "passA",
            "userB": "passB",
        }

        writer.clear_passwords()

        assert writer._thread_local.passwords == {}

    def test_clear_passwords_is_noop_when_no_passwords_set(self):
        from local_deep_research.database.thread_metrics import (
            ThreadSafeMetricsWriter,
        )

        writer = ThreadSafeMetricsWriter()
        # No passwords dict exists; clear_passwords should not raise.
        writer.clear_passwords()
        assert not hasattr(writer._thread_local, "passwords")

    def test_cleanup_current_thread_clears_passwords(self):
        """Regression test — cleanup_current_thread() must clear the passwords dict
        on the module-level metrics_writer so pooled workers don't retain creds."""
        from local_deep_research.database.thread_metrics import metrics_writer
        from local_deep_research.database.thread_local_session import (
            cleanup_current_thread,
        )

        metrics_writer.set_user_password("leaked_user", "leaked_pass")
        assert (
            metrics_writer._thread_local.passwords.get("leaked_user")
            == "leaked_pass"
        )

        cleanup_current_thread()

        assert metrics_writer._thread_local.passwords == {}


class TestThreadIsolation:
    """Tests for thread isolation of password storage."""

    def test_passwords_are_thread_isolated(self):
        """Test that passwords stored in one thread are not visible in another."""
        from local_deep_research.database.thread_metrics import (
            ThreadSafeMetricsWriter,
        )

        writer = ThreadSafeMetricsWriter()
        results = {}

        def thread1_work():
            writer.set_user_password("user1", "pass1")
            results["thread1_has_password"] = hasattr(
                writer._thread_local, "passwords"
            )
            results["thread1_user1"] = writer._thread_local.passwords.get(
                "user1"
            )

        def thread2_work():
            # Small delay to ensure thread1 runs first
            import time

            time.sleep(0.01)
            results["thread2_has_password"] = hasattr(
                writer._thread_local, "passwords"
            )

        t1 = threading.Thread(target=thread1_work)
        t2 = threading.Thread(target=thread2_work)

        t1.start()
        t2.start()
        t1.join()
        t2.join()

        # Thread 1 should have its password
        assert results["thread1_has_password"] is True
        assert results["thread1_user1"] == "pass1"

        # Thread 2 should NOT have access to thread 1's passwords
        assert results["thread2_has_password"] is False