learningcircuit--local-deep-research/tests/chat/test_chat_concurrency.py

"""Unit tests for ChatService concurrency safety.

These tests verify thread-safety for concurrent operations.
Uses concurrent.futures.ThreadPoolExecutor to simulate concurrent access.
"""

import uuid
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import UTC, datetime
from unittest.mock import MagicMock, patch
from contextlib import contextmanager
import threading

from tests.chat.conftest import setup_query_mock_with_session


class TestChatConcurrency:
    """Tests for concurrent operations on ChatService."""

    def test_concurrent_session_creation_unique_ids(self):
        """Test that concurrent session creation generates unique IDs."""
        from src.local_deep_research.chat.service import ChatService

        created_ids = []
        lock = threading.Lock()

        @contextmanager
        def mock_get_user_db_session(username, password=None):
            mock_session = MagicMock()
            mock_session.add = MagicMock()
            mock_session.commit = MagicMock()
            yield mock_session

        def create_session_task():
            service = ChatService(username="testuser")
            session_id = service.create_session(
                initial_query=f"Query {threading.current_thread().name}"
            )
            with lock:
                created_ids.append(session_id)
            return session_id

        # Apply patch at test level (outside threads) to avoid race conditions
        # during patch application/removal across threads
        with patch(
            "src.local_deep_research.chat.service.get_user_db_session",
            mock_get_user_db_session,
        ):
            # Create sessions concurrently
            num_threads = 10
            with ThreadPoolExecutor(max_workers=num_threads) as executor:
                futures = [
                    executor.submit(create_session_task)
                    for _ in range(num_threads)
                ]
                for future in as_completed(futures):
                    future.result()  # Ensure all complete

        # All IDs should be unique
        assert len(created_ids) == num_threads
        assert len(set(created_ids)) == num_threads  # All unique
        # All should be valid UUIDs
        for session_id in created_ids:
            uuid.UUID(session_id)  # Raises ValueError if invalid

    def test_multiple_sessions_same_user_isolated(self):
        """Test that multiple sessions for the same user are independent."""
        from src.local_deep_research.chat.service import ChatService

        sessions_created = []
        lock = threading.Lock()

        # Track message counts for each session
        session_message_counts = {}

        @contextmanager
        def mock_get_user_db_session(username, password=None):
            mock_session = MagicMock()
            mock_session.add = MagicMock()
            mock_session.commit = MagicMock()

            # Mock session lookup to return appropriate mock
            def mock_filter_by(id=None, **kwargs):
                result = MagicMock()
                if id in session_message_counts:
                    mock_chat_session = MagicMock()
                    mock_chat_session.id = id
                    mock_chat_session.message_count = session_message_counts[id]
                    result.first.return_value = mock_chat_session
                else:
                    result.first.return_value = None
                return result

            mock_session.query.return_value.filter_by = mock_filter_by
            yield mock_session

        def create_and_use_session():
            service = ChatService(username="testuser")
            session_id = service.create_session()

            with lock:
                sessions_created.append(session_id)
                session_message_counts[session_id] = 0

            return session_id

        # Patch at module level before spawning threads
        with patch(
            "src.local_deep_research.chat.service.get_user_db_session",
            mock_get_user_db_session,
        ):
            # Create multiple sessions concurrently
            num_sessions = 5
            with ThreadPoolExecutor(max_workers=num_sessions) as executor:
                futures = [
                    executor.submit(create_and_use_session)
                    for _ in range(num_sessions)
                ]
                for f in as_completed(futures):
                    f.result()

        # Each session should be created independently
        assert len(sessions_created) == num_sessions
        assert len(set(sessions_created)) == num_sessions

    def test_concurrent_reads_during_write_consistent(self):
        """Test that reads during writes return consistent data."""
        from src.local_deep_research.chat.service import ChatService

        # Shared state for the mock
        mock_sessions = {}
        write_lock = threading.Lock()
        read_results = []
        results_lock = threading.Lock()

        @contextmanager
        def mock_get_user_db_session(username, password=None):
            mock_session = MagicMock()
            mock_session.add = MagicMock()
            mock_session.commit = MagicMock()

            def mock_filter_by(id=None, status="active", **kwargs):
                result = MagicMock()
                with write_lock:
                    if id and id in mock_sessions:
                        session_obj = MagicMock()
                        session_obj.id = id
                        session_obj.title = mock_sessions[id]["title"]
                        session_obj.status = "active"
                        session_obj.message_count = 0
                        session_obj.created_at = datetime(
                            2024, 1, 1, tzinfo=UTC
                        )
                        result.first.return_value = session_obj
                    else:
                        result.first.return_value = None

                    # For list queries
                    sessions_list = [
                        MagicMock(
                            id=sid,
                            title=sdata["title"],
                            status="active",
                            message_count=0,
                            created_at=datetime(2024, 1, 1, tzinfo=UTC),
                        )
                        for sid, sdata in mock_sessions.items()
                    ]
                    result.order_by.return_value.offset.return_value.limit.return_value.all.return_value = sessions_list

                return result

            mock_session.query.return_value.filter_by = mock_filter_by
            yield mock_session

        def write_task(session_id, title):
            """Write task that creates/updates a session."""
            with write_lock:
                mock_sessions[session_id] = {"title": title}

        def read_task():
            """Read task that lists sessions."""
            service = ChatService(username="testuser")
            sessions = service.list_sessions()
            with results_lock:
                read_results.append(len(sessions))
            return sessions

        # First create some sessions
        for i in range(5):
            write_task(f"session-{i}", f"Title {i}")

        # Apply patch at test level (outside threads) to avoid race conditions
        with patch(
            "src.local_deep_research.chat.service.get_user_db_session",
            mock_get_user_db_session,
        ):
            # Now do concurrent reads
            with ThreadPoolExecutor(max_workers=10) as executor:
                futures = [executor.submit(read_task) for _ in range(10)]
                for future in as_completed(futures):
                    future.result()

        # All reads should have gotten consistent results
        assert len(read_results) == 10
        # All should have found the same number of sessions (5)
        assert all(count == 5 for count in read_results)

    def test_rapid_message_sends_no_data_loss(self):
        """Test that rapidly sending messages doesn't lose any messages."""
        from src.local_deep_research.chat.service import ChatService

        messages_added = []
        lock = threading.Lock()

        mock_session_obj = MagicMock()
        mock_session_obj.id = "test-session"
        mock_session_obj.message_count = 0

        @contextmanager
        def mock_get_user_db_session(username, password=None):
            mock_session = MagicMock()

            def mock_add(message):
                with lock:
                    messages_added.append(message.content)

            mock_session.add = mock_add
            mock_session.commit = MagicMock()
            # Set up query mock to handle with_for_update() chain
            setup_query_mock_with_session(mock_session, mock_session_obj)
            yield mock_session

        def send_message_task(msg_content):
            service = ChatService(username="testuser")
            message_id = service.add_message(
                session_id="test-session",
                role="user",
                content=msg_content,
                message_type="query",
            )
            return message_id

        # Patch at module level before spawning threads
        with patch(
            "src.local_deep_research.chat.service.get_user_db_session",
            mock_get_user_db_session,
        ):
            # Send many messages rapidly
            num_messages = 20
            message_contents = [f"Message {i}" for i in range(num_messages)]

            with ThreadPoolExecutor(max_workers=10) as executor:
                futures = [
                    executor.submit(send_message_task, content)
                    for content in message_contents
                ]
                results = [f.result() for f in as_completed(futures)]

        # All messages should have been sent
        assert len(messages_added) == num_messages
        # All message IDs should be returned
        assert len(results) == num_messages
        # All message contents should be present (order may vary)
        assert set(messages_added) == set(message_contents)

    def test_concurrent_context_updates_preserve_data(self):
        """Test that concurrent context updates don't lose entities."""
        from src.local_deep_research.chat.service import ChatService

        # The implementation uses with_for_update() for row-level locking
        # to ensure atomic read-modify-write operations

        mock_session_obj = MagicMock()
        mock_session_obj.id = "test-session"
        mock_session_obj.accumulated_context = {
            "key_entities": [],
            "topics": [],
            "summary": "",
        }

        lock = threading.Lock()
        update_count = 0

        @contextmanager
        def mock_get_user_db_session(username, password=None):
            nonlocal update_count
            # with_for_update() is a no-op on SQLite (see the service comment),
            # so the read-modify-write of accumulated_context is not serialized
            # by the DB. In production, concurrent updates to a SINGLE session's
            # context do not occur because only one research runs per chat
            # session at a time (the per-session in-flight guard). Hold `lock`
            # for the whole session context here to mirror that real-world
            # serialization, so the test deterministically verifies the MERGE
            # accumulates every entity rather than clobbering. (Full
            # multi-writer fidelity would require a real DB.)
            with lock:
                mock_session = MagicMock()
                # Set up query mock to handle with_for_update() chain
                setup_query_mock_with_session(mock_session, mock_session_obj)

                def mock_commit():
                    nonlocal update_count
                    update_count += 1

                mock_session.commit = mock_commit
                yield mock_session

        def update_context_task(entity_name):
            service = ChatService(username="testuser")
            service.update_accumulated_context(
                "test-session",
                new_entities=[entity_name],
            )
            return entity_name

        # Update context concurrently
        num_updates = 5
        entity_names = [f"entity_{i}" for i in range(num_updates)]

        # Apply patch at test level (outside threads) to avoid race conditions
        with patch(
            "src.local_deep_research.chat.service.get_user_db_session",
            mock_get_user_db_session,
        ):
            with ThreadPoolExecutor(max_workers=num_updates) as executor:
                futures = [
                    executor.submit(update_context_task, name)
                    for name in entity_names
                ]
                for future in as_completed(futures):
                    future.result()

        # All updates should have been committed
        assert update_count == num_updates
        # And — the actual point of the test — the merge must have ACCUMULATED
        # every entity rather than clobbering prior ones: a regression that
        # overwrote key_entities on each update would still pass a commit count.
        preserved = mock_session_obj.accumulated_context["key_entities"]
        assert set(preserved) == set(entity_names)

    def test_thread_local_services_independent(self):
        """Test that ChatService instances in different threads are independent."""
        from src.local_deep_research.chat.service import ChatService

        service_ids = []
        lock = threading.Lock()

        def create_service_task():
            service = ChatService(
                username=f"user_{threading.current_thread().name}"
            )
            with lock:
                service_ids.append(id(service))
            return service

        # Create services in different threads
        num_threads = 5
        with ThreadPoolExecutor(max_workers=num_threads) as executor:
            futures = [
                executor.submit(create_service_task) for _ in range(num_threads)
            ]
            services = [f.result() for f in as_completed(futures)]

        # Each thread should have its own service instance
        assert len(services) == num_threads
        assert len(set(service_ids)) == num_threads  # All unique object IDs


class TestAddMessageSequenceRace:
    """Real-SQLite (not mocked) coverage for the add_message sequence race.

    test_rapid_message_sends_no_data_loss mocks the DB entirely and
    therefore cannot catch IntegrityError collisions on
    uq_chat_message_session_seq. These tests exercise the actual retry path
    against a file-backed SQLite DB.
    """

    def _setup_real_db(self, tmp_path):
        """Create a file-backed SQLite DB with chat tables + a seeded session.

        :memory: isn't usable here because it isn't shared across threads.
        """
        import uuid
        from datetime import datetime, UTC
        from sqlalchemy import create_engine
        from sqlalchemy.orm import sessionmaker
        from src.local_deep_research.database.models import (
            Base,
            ChatSession,
        )

        db_path = tmp_path / "chat_race.db"
        engine = create_engine(f"sqlite:///{db_path}")
        Base.metadata.create_all(engine)

        Session = sessionmaker(bind=engine)
        session_id = str(uuid.uuid4())
        with Session() as db:
            db.add(
                ChatSession(
                    id=session_id,
                    title="race test",
                    status="active",
                    accumulated_context={
                        "key_entities": [],
                        "topics": [],
                        "summary": "",
                    },
                    created_at=datetime.now(UTC),
                    message_count=0,
                )
            )
            db.commit()

        return engine, Session, session_id

    def test_concurrent_add_message_unique_sequence_numbers(self, tmp_path):
        """10 concurrent add_message calls produce 10 unique sequence_numbers.

        This exercises the add_message sequence-number race: without
        retry-on-IntegrityError, two callers would both compute the same
        `session.message_count + 1`
        and one would fail with a 500 when the uq_chat_message_session_seq
        constraint fires.
        """
        from contextlib import contextmanager
        from src.local_deep_research.chat.service import ChatService
        from src.local_deep_research.database.models import ChatMessage

        engine, Session, session_id = self._setup_real_db(tmp_path)

        @contextmanager
        def real_get_user_db_session(username, password=None):
            with Session() as db:
                yield db

        service = ChatService(username="testuser")
        num_messages = 10
        results = []
        errors = []
        errors_lock = threading.Lock()

        def add_one(i):
            try:
                return service.add_message(
                    session_id=session_id,
                    role="user",
                    content=f"message {i}",
                    message_type="query",
                )
            except Exception as e:  # noqa: BLE001
                with errors_lock:
                    errors.append(e)
                return None

        with patch(
            "src.local_deep_research.chat.service.get_user_db_session",
            real_get_user_db_session,
        ):
            with ThreadPoolExecutor(max_workers=num_messages) as executor:
                futures = [
                    executor.submit(add_one, i) for i in range(num_messages)
                ]
                for f in as_completed(futures):
                    results.append(f.result())

        assert errors == [], f"add_message surfaced errors: {errors}"
        assert len([r for r in results if r]) == num_messages

        with Session() as db:
            sequences = sorted(
                r[0]
                for r in db.query(ChatMessage.sequence_number)
                .filter_by(session_id=session_id)
                .all()
            )
        # 10 unique sequence numbers, 1..10, no collisions
        assert sequences == list(range(1, num_messages + 1))
        engine.dispose()