"""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