LangBot/tests/unit_tests/rag/test_file_storage.py

"""Unit tests for RuntimeKnowledgeBase file storage behavior."""

from __future__ import annotations

import contextvars
import io
import zipfile
from contextlib import asynccontextmanager
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock

import pytest

from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.core.taskmgr import TaskCapacityError
from langbot.pkg.rag.knowledge.kbmgr import RuntimeKnowledgeBase
from langbot.pkg.storage.mgr import StorageMgr
from langbot.pkg.workspace.errors import WorkspaceNotFoundError


WORKSPACE_A = '00000000-0000-0000-0000-00000000000a'
CONTEXT = ExecutionContext(
    instance_uuid='instance-a',
    workspace_uuid=WORKSPACE_A,
    placement_generation=2,
)


def _upload_key(logical_key: str, *, context: ExecutionContext = CONTEXT) -> str:
    return StorageMgr.scoped_object_key(
        context,
        owner_type='upload_document',
        owner='account:test',
        key=logical_key,
    )


def _make_zip_bytes(entries: dict[str, bytes], *, compression: int = zipfile.ZIP_STORED) -> bytes:
    buffer = io.BytesIO()
    with zipfile.ZipFile(buffer, 'w', compression=compression) as zf:
        for name, content in entries.items():
            zf.writestr(name, content)
        zf.mkdir('emptydir')
    return buffer.getvalue()


def _make_app() -> Mock:
    app = Mock()
    app.logger = Mock()
    app.task_mgr = Mock()
    storage_mgr = StorageMgr(app)
    storage_mgr.storage_provider = Mock()
    storage_mgr.storage_provider.exists = AsyncMock(return_value=True)
    storage_mgr.storage_provider.load = AsyncMock()
    storage_mgr.storage_provider.load_bounded = AsyncMock()
    storage_mgr.storage_provider.save = AsyncMock()
    storage_mgr.storage_provider.size = AsyncMock(return_value=123)
    storage_mgr.storage_provider.delete = AsyncMock()
    app.storage_mgr = storage_mgr
    app.persistence_mgr = Mock()
    app.persistence_mgr.execute_async = AsyncMock(return_value=SimpleNamespace(rowcount=1))
    app.plugin_connector = Mock()
    app.plugin_connector.require_workspace_context = AsyncMock(side_effect=lambda context: context)
    app.workspace_service = SimpleNamespace(
        get_execution_binding=AsyncMock(
            return_value=SimpleNamespace(
                instance_uuid=CONTEXT.instance_uuid,
                workspace_uuid=CONTEXT.workspace_uuid,
                placement_generation=CONTEXT.placement_generation,
            )
        )
    )
    return app


def _make_kb(plugin_id: str | None = 'author/engine') -> RuntimeKnowledgeBase:
    kb_entity = Mock()
    kb_entity.uuid = 'test-kb-uuid'
    kb_entity.workspace_uuid = WORKSPACE_A
    kb_entity.collection_id = 'test-collection'
    kb_entity.creation_settings = {}
    kb_entity.knowledge_engine_plugin_id = plugin_id
    return RuntimeKnowledgeBase(_make_app(), kb_entity, CONTEXT)


class TestStoreFile:
    @pytest.mark.asyncio
    async def test_store_file_creates_pending_record_and_user_task(self):
        kb = _make_kb()

        def create_user_task(coro, **kwargs):
            coro.close()
            return SimpleNamespace(id='task-1', kwargs=kwargs)

        kb.ap.task_mgr.create_user_task = Mock(side_effect=create_user_task)

        object_key = _upload_key('documents/test.pdf')
        task_id = await kb.store_file(CONTEXT, object_key)

        assert task_id == 'task-1'
        kb.ap.storage_mgr.storage_provider.exists.assert_awaited_once_with(object_key)
        kb.ap.persistence_mgr.execute_async.assert_awaited_once()
        call_kwargs = kb.ap.task_mgr.create_user_task.call_args.kwargs
        assert call_kwargs['kind'] == 'knowledge-operation'
        assert call_kwargs['name'] == f'knowledge-store-file-{object_key}'
        assert call_kwargs['label'] == f'Store file {object_key}'

    @pytest.mark.asyncio
    async def test_store_file_raises_when_source_file_missing(self):
        kb = _make_kb()
        kb.ap.storage_mgr.storage_provider.exists = AsyncMock(return_value=False)

        object_key = _upload_key('missing.pdf')
        with pytest.raises(WorkspaceNotFoundError, match='Upload not found'):
            await kb.store_file(CONTEXT, object_key)

        kb.ap.persistence_mgr.execute_async.assert_not_awaited()
        kb.ap.task_mgr.create_user_task.assert_not_called()

    @pytest.mark.asyncio
    async def test_store_file_rejects_cross_workspace_upload_key(self):
        kb = _make_kb()
        other_context = ExecutionContext(
            instance_uuid=CONTEXT.instance_uuid,
            workspace_uuid='00000000-0000-0000-0000-00000000000b',
            placement_generation=CONTEXT.placement_generation,
        )

        with pytest.raises(WorkspaceNotFoundError, match='Upload not found'):
            await kb.store_file(CONTEXT, _upload_key('stolen.pdf', context=other_context))

        kb.ap.storage_mgr.storage_provider.exists.assert_not_awaited()
        kb.ap.persistence_mgr.execute_async.assert_not_awaited()

    @pytest.mark.asyncio
    async def test_store_file_rejects_stale_generation_and_wrong_owner_type(self):
        kb = _make_kb()
        stale_context = ExecutionContext(
            instance_uuid=CONTEXT.instance_uuid,
            workspace_uuid=CONTEXT.workspace_uuid,
            placement_generation=CONTEXT.placement_generation + 1,
        )
        plugin_key = StorageMgr.scoped_object_key(
            CONTEXT,
            owner_type='plugin_config',
            owner='plugin:test',
            key='config.pdf',
        )

        for object_key in (_upload_key('stale.pdf', context=stale_context), plugin_key, 'raw.pdf'):
            with pytest.raises(WorkspaceNotFoundError, match='Upload not found'):
                await kb.store_file(CONTEXT, object_key)

        kb.ap.storage_mgr.storage_provider.exists.assert_not_awaited()

    @pytest.mark.asyncio
    async def test_store_file_rolls_back_pending_record_when_task_capacity_is_exhausted(self):
        kb = _make_kb()
        object_key = _upload_key('queued.pdf')

        def reject(coro, **_kwargs):
            coro.close()
            raise TaskCapacityError('capacity')

        kb.ap.task_mgr.create_user_task.side_effect = reject

        with pytest.raises(TaskCapacityError, match='capacity'):
            await kb.store_file(CONTEXT, object_key)

        statements = [str(call.args[0]) for call in kb.ap.persistence_mgr.execute_async.await_args_list]
        assert any(statement.startswith('INSERT') for statement in statements)
        assert any(statement.startswith('DELETE') for statement in statements)


class TestStoreZipFile:
    @pytest.mark.asyncio
    async def test_store_zip_file_extracts_supported_files_and_skips_noise(self):
        kb = _make_kb()
        kb.ap.storage_mgr.storage_provider.load_bounded = AsyncMock(
            return_value=_make_zip_bytes(
                {
                    'doc1.pdf': b'pdf',
                    'doc2.txt': b'text',
                    'subdir/doc3.md': b'markdown',
                    'page.html': b'html',
                    'image.png': b'png',
                    '.hidden': b'hidden',
                    '__MACOSX/doc1.pdf': b'metadata',
                }
            )
        )
        kb.store_file = AsyncMock(side_effect=['task-pdf', 'task-txt', 'task-md', 'task-html'])

        zip_key = _upload_key('archive.zip')
        task_id = await kb._store_zip_file(CONTEXT, zip_key, parser_plugin_id='parser/plugin')

        assert task_id == 'task-pdf'
        assert kb.ap.storage_mgr.storage_provider.save.await_count == 4
        saved_names = [call.args[0] for call in kb.ap.storage_mgr.storage_provider.save.await_args_list]
        assert {name.rsplit('.', 1)[-1] for name in saved_names} == {'pdf', 'txt', 'md', 'html'}
        for name in saved_names:
            StorageMgr.require_scoped_object_key(
                CONTEXT,
                name,
                expected_owner_type='upload_document',
            )
        forwarded_keys = [call.args[1] for call in kb.store_file.await_args_list]
        assert forwarded_keys == saved_names
        kb.ap.storage_mgr.storage_provider.delete.assert_awaited_once_with(zip_key)

    @pytest.mark.asyncio
    async def test_store_zip_file_raises_when_no_supported_files(self):
        kb = _make_kb()
        kb.ap.storage_mgr.storage_provider.load_bounded = AsyncMock(
            return_value=_make_zip_bytes({'image.png': b'png', 'video.mp4': b'video'})
        )
        kb.store_file = AsyncMock()

        with pytest.raises(Exception, match='No supported files found'):
            await kb._store_zip_file(CONTEXT, _upload_key('archive.zip'))

        kb.store_file.assert_not_awaited()
        kb.ap.storage_mgr.storage_provider.delete.assert_awaited_once_with(_upload_key('archive.zip'))

    @pytest.mark.asyncio
    async def test_store_zip_file_rejects_too_many_documents_before_extracting(self):
        kb = _make_kb()
        kb.ap.storage_mgr.storage_provider.load_bounded = AsyncMock(
            return_value=_make_zip_bytes({f'doc-{index}.txt': b'text' for index in range(9)})
        )
        kb.store_file = AsyncMock()

        with pytest.raises(ValueError, match='too many supported documents'):
            await kb._store_zip_file(CONTEXT, _upload_key('archive.zip'))

        kb.store_file.assert_not_awaited()
        kb.ap.storage_mgr.storage_provider.save.assert_not_awaited()
        kb.ap.storage_mgr.storage_provider.delete.assert_awaited_once_with(_upload_key('archive.zip'))

    @pytest.mark.asyncio
    async def test_store_zip_file_rejects_extreme_compression_ratio_before_extracting(self):
        kb = _make_kb()
        kb.ap.storage_mgr.storage_provider.load_bounded = AsyncMock(
            return_value=_make_zip_bytes(
                {'bomb.txt': b'A' * (1024 * 1024)},
                compression=zipfile.ZIP_DEFLATED,
            )
        )
        kb.store_file = AsyncMock()

        with pytest.raises(ValueError, match='compression-ratio limit'):
            await kb._store_zip_file(CONTEXT, _upload_key('archive.zip'))

        kb.store_file.assert_not_awaited()
        kb.ap.storage_mgr.storage_provider.save.assert_not_awaited()
        kb.ap.storage_mgr.storage_provider.delete.assert_awaited_once_with(_upload_key('archive.zip'))


class TestStoreFileTask:
    @pytest.mark.asyncio
    async def test_store_file_task_opens_uow_before_first_database_helper(self):
        kb = _make_kb()
        active_workspace = contextvars.ContextVar('rag_task_workspace', default=None)
        observed = []

        @asynccontextmanager
        async def tenant_uow(workspace_uuid):
            token = active_workspace.set(workspace_uuid)
            try:
                yield
            finally:
                active_workspace.reset(token)

        kb.ap.persistence_mgr.mode = SimpleNamespace(value='cloud_runtime')
        kb.ap.persistence_mgr.tenant_uow = tenant_uow

        async def assert_execution_context(_context):
            observed.append(active_workspace.get())

        kb._assert_execution_context = AsyncMock(side_effect=assert_execution_context)
        kb._set_file_status = AsyncMock(side_effect=[True, True])
        kb._ingest_document = AsyncMock(return_value={'status': 'completed'})
        object_key = _upload_key('scoped.pdf')
        file_obj = SimpleNamespace(uuid='file-uuid', file_name=object_key, extension='pdf')

        await kb._store_file_task(CONTEXT, file_obj, Mock())

        assert observed[0] == WORKSPACE_A

    @pytest.mark.asyncio
    async def test_store_file_task_marks_completed_and_cleans_storage(self):
        kb = _make_kb()
        kb._ingest_document = AsyncMock(return_value={'status': 'completed'})
        object_key = _upload_key('test.pdf')
        file_obj = SimpleNamespace(uuid='file-uuid', file_name=object_key, extension='pdf')
        task_context = Mock()

        await kb._store_file_task(CONTEXT, file_obj, task_context)

        task_context.set_current_action.assert_called_once_with('Processing file')
        kb.ap.storage_mgr.storage_provider.size.assert_awaited_once_with(object_key)
        kb._ingest_document.assert_awaited_once()
        assert kb.ap.persistence_mgr.execute_async.await_count == 2
        kb.ap.storage_mgr.storage_provider.delete.assert_awaited_once_with(object_key)

    @pytest.mark.asyncio
    async def test_store_file_task_marks_failed_and_cleans_storage(self):
        kb = _make_kb()
        kb._ingest_document = AsyncMock(return_value={'status': 'failed', 'error_message': 'parser failed'})
        object_key = _upload_key('bad.pdf')
        file_obj = SimpleNamespace(uuid='file-uuid', file_name=object_key, extension='pdf')
        task_context = Mock()

        with pytest.raises(Exception, match='parser failed'):
            await kb._store_file_task(CONTEXT, file_obj, task_context)

        assert kb.ap.persistence_mgr.execute_async.await_count == 2
        kb.ap.storage_mgr.storage_provider.delete.assert_awaited_once_with(object_key)


class TestDeleteDocument:
    @pytest.mark.asyncio
    async def test_delete_document_returns_false_when_no_plugin_id(self):
        kb = _make_kb(plugin_id=None)

        result = await kb._delete_document(CONTEXT, 'doc-id')

        assert result is False

    @pytest.mark.asyncio
    async def test_delete_document_calls_configured_rag_plugin(self):
        kb = _make_kb()
        kb.ap.plugin_connector.call_rag_delete_document = AsyncMock(return_value=True)

        result = await kb._delete_document(CONTEXT, 'doc-id')

        assert result is True
        kb.ap.plugin_connector.call_rag_delete_document.assert_awaited_once_with(
            'author/engine', 'doc-id', 'test-kb-uuid'
        )

    @pytest.mark.asyncio
    async def test_delete_document_returns_false_on_plugin_error(self):
        kb = _make_kb()
        kb.ap.plugin_connector.call_rag_delete_document = AsyncMock(side_effect=Exception('plugin error'))

        result = await kb._delete_document(CONTEXT, 'doc-id')

        assert result is False
        kb.ap.logger.error.assert_called_once()