diff --git a/tests/unit_tests/agent/test_chat_handler.py b/tests/unit_tests/agent/test_chat_handler.py index 6eede3f28..10ac2109e 100644 --- a/tests/unit_tests/agent/test_chat_handler.py +++ b/tests/unit_tests/agent/test_chat_handler.py @@ -580,3 +580,4 @@ class TestChatHandlerAsyncBehavior: assert len(results) == 1 assert results[0].result_type == entities.ResultType.CONTINUE assert len(query.resp_messages) == 1 + assert query.resp_messages[0].content == 'Reply from plugin' diff --git a/tests/unit_tests/agent/test_event_first_protocol.py b/tests/unit_tests/agent/test_event_first_protocol.py index ec41da89e..2721ff630 100644 --- a/tests/unit_tests/agent/test_event_first_protocol.py +++ b/tests/unit_tests/agent/test_event_first_protocol.py @@ -303,6 +303,10 @@ class TestRunnerContextProtocolV1: assert ctx.event is not None assert ctx.event.event_type == 'message.received' + from pydantic import ValidationError + + with pytest.raises(ValidationError, match='event'): + RunnerContext.model_validate(ctx.model_dump(exclude={'event'})) def test_sdk_context_has_no_history_message_fields(self): """RunnerContext should not expose inline history message fields.""" @@ -345,6 +349,8 @@ class TestSDKResultProtocolV1: ) assert result.run_id == 'run_1' + with pytest.raises(TypeError, match='run_id'): + RunnerResult.message_completed(message=Message(role='assistant', content='Hello')) # Fixtures diff --git a/tests/unit_tests/agent/test_event_log_transcript.py b/tests/unit_tests/agent/test_event_log_transcript.py index 28540071c..c31468751 100644 --- a/tests/unit_tests/agent/test_event_log_transcript.py +++ b/tests/unit_tests/agent/test_event_log_transcript.py @@ -139,30 +139,31 @@ class TestTranscriptStore: """Test TranscriptStore operations.""" @pytest.mark.asyncio - async def test_append_transcript(self, mock_db_engine): + async def test_append_transcript(self): """Test appending a transcript item.""" - from unittest.mock import AsyncMock, MagicMock, patch + from sqlalchemy.ext.asyncio import create_async_engine + from langbot.pkg.entity.persistence.base import Base + from uuid import UUID - store = TranscriptStore(mock_db_engine) - - mock_session = AsyncMock() - mock_session.add = MagicMock() - mock_session.commit = AsyncMock() - - # Mock _get_next_seq - with patch.object(store, '_get_next_seq', return_value=1): - with patch.object(store, '_session_factory') as mock_factory: - mock_factory.return_value.__aenter__.return_value = mock_session - - transcript_id = await store.append_transcript( - transcript_id=None, # Auto-generate - event_id='evt_1', - conversation_id='conv_1', - role='user', - content='Hello', - ) - - assert transcript_id is not None + engine = create_async_engine('sqlite+aiosqlite:///:memory:') + try: + async with engine.begin() as connection: + await connection.run_sync(Base.metadata.create_all) + store = TranscriptStore(engine) + transcript_id = await store.append_transcript( + transcript_id=None, + event_id='evt_1', + conversation_id='conv_1', + role='user', + content='Hello', + ) + assert UUID(transcript_id).version == 4 + items, _, _, _ = await store.page_transcript('conv_1') + assert [(item['transcript_id'], item['event_id'], item['role'], item['content']) for item in items] == [ + (transcript_id, 'evt_1', 'user', 'Hello') + ] + finally: + await engine.dispose() class TestEventLogStoreRealSQLite: @@ -258,6 +259,7 @@ class TestEventLogStoreRealSQLite: async def test_get_latest_cursor(self, db_engine): """Test get_latest_cursor with real DB.""" store = EventLogStore(db_engine) + assert await store.get_latest_cursor('conv_cursor') is None # Append events for i in range(3): @@ -270,8 +272,12 @@ class TestEventLogStoreRealSQLite: # Get latest cursor cursor = await store.get_latest_cursor('conv_cursor') - assert cursor is not None - assert int(cursor) > 0 + assert cursor == '3' + await store.append_event( + event_id='other', event_type='message.received', source='platform', conversation_id='other' + ) + assert await store.get_latest_cursor('conv_cursor') == '3' + assert await store.get_latest_cursor('other') == '4' @pytest.mark.asyncio async def test_cleanup_events_older_than(self, db_engine): diff --git a/tests/unit_tests/agent/test_orchestrator_integration.py b/tests/unit_tests/agent/test_orchestrator_integration.py index 61802e15e..fa5fb18d1 100644 --- a/tests/unit_tests/agent/test_orchestrator_integration.py +++ b/tests/unit_tests/agent/test_orchestrator_integration.py @@ -789,6 +789,13 @@ async def test_orchestrator_applies_state_updates_and_suppresses_protocol_event( assert [message.content for message in messages] == ['state saved'] # State is persisted to the database via PersistentStateStore. + from langbot.pkg.agent.runner.persistent_state_store import get_persistent_state_store + + event = QueryEntryAdapter.query_to_event(query) + config = QueryEntryAdapter.config_to_agent_config(query, RUNNER_ID) + binding = AgentBindingResolver().resolve_one(event, [config]) + snapshot = await get_persistent_state_store(db_engine).build_snapshot_from_event(event, binding, descriptor) + assert snapshot['conversation'].get('external.conversation_id') == 'external_conv_123' @pytest.mark.asyncio diff --git a/tests/unit_tests/agent/test_registry.py b/tests/unit_tests/agent/test_registry.py index da7780622..025d23f74 100644 --- a/tests/unit_tests/agent/test_registry.py +++ b/tests/unit_tests/agent/test_registry.py @@ -124,6 +124,7 @@ class TestRegistryDiscovery: async def test_discover_caches_results(self): """Discovery should cache results.""" ap = FakeApplication() + ap.plugin_connector.list_runners = AsyncMock(wraps=ap.plugin_connector.list_runners) registry = RunnerRegistry(ap) # First discovery @@ -134,6 +135,11 @@ class TestRegistryDiscovery: assert registry._cache is not None assert len(runners1) == len(runners2) + ap.plugin_connector.list_runners.assert_awaited_once() + assert [runner.id for runner in runners2] == [ + 'plugin:langbot-team/LocalAgent/default', + 'plugin:alice/my-agent/custom', + ] @pytest.mark.asyncio async def test_discover_handles_plugin_disabled(self): diff --git a/tests/unit_tests/agent/test_result_normalizer.py b/tests/unit_tests/agent/test_result_normalizer.py index b4dbe0970..ca1bd0e39 100644 --- a/tests/unit_tests/agent/test_result_normalizer.py +++ b/tests/unit_tests/agent/test_result_normalizer.py @@ -156,6 +156,7 @@ class TestNormalizeRunCompleted: assert result is not None assert isinstance(result, provider_message.Message) + assert (result.role, result.content) == ('assistant', 'Final response') @pytest.mark.asyncio async def test_normalize_run_completed_without_message(self): diff --git a/tests/unit_tests/agent/test_run_ledger_api_auth.py b/tests/unit_tests/agent/test_run_ledger_api_auth.py index 992bbcdfd..e75f67739 100644 --- a/tests/unit_tests/agent/test_run_ledger_api_auth.py +++ b/tests/unit_tests/agent/test_run_ledger_api_auth.py @@ -173,6 +173,22 @@ async def _create_run( ) +@pytest.fixture +async def ledger_snapshot(db_engine): + await _create_run( + db_engine, run_id='workspace_guard', workspace_id='workspace_guard', plugin_identity='guard/runner' + ) + + async def snapshot(): + async with AsyncSession(db_engine) as session: + return [ + (await session.execute(sqlalchemy.select(*model.__table__.c).order_by(model.id))).mappings().all() + for model in (agent_run_model.AgentRun, agent_run_model.AgentRunEvent) + ] + + return snapshot + + @pytest.mark.asyncio async def test_run_get_returns_current_run(session_registry, db_engine): await _register_session(session_registry, available_apis={'run_get': True}) @@ -262,6 +278,7 @@ async def test_run_list_filters_same_scope_different_runner_owner(session_regist async def test_non_admin_target_run_actions_reject_same_scope_different_runner_owner( session_registry, db_engine, + ledger_snapshot, ): await _register_session( session_registry, @@ -349,10 +366,12 @@ async def test_non_admin_target_run_actions_reject_same_scope_different_runner_o ), ] + before = await ledger_snapshot() for action, payload in calls: result = await handler.actions[action](payload) assert result.code != 0, action assert 'not accessible' in result.message.lower() + assert await ledger_snapshot() == before @pytest.mark.asyncio @@ -1046,11 +1065,12 @@ async def test_run_append_result_basic_path(session_registry, db_engine): @pytest.mark.asyncio -async def test_run_append_result_rejects_side_effecting_result_type(session_registry, db_engine): +async def test_run_append_result_rejects_side_effecting_result_type(session_registry, db_engine, ledger_snapshot): await _register_session(session_registry, available_apis={'run_append_result': True}) await _create_run(db_engine) handler = _handler(db_engine) run_append_result = handler.actions[PluginToRuntimeAction.RUN_APPEND_RESULT.value] + before = await ledger_snapshot() result = await run_append_result( { @@ -1066,14 +1086,16 @@ async def test_run_append_result_rejects_side_effecting_result_type(session_regi assert result.code != 0 assert 'canonical runner result path' in result.message + assert await ledger_snapshot() == before @pytest.mark.asyncio -async def test_run_append_result_rejects_unknown_result_type(session_registry, db_engine): +async def test_run_append_result_rejects_unknown_result_type(session_registry, db_engine, ledger_snapshot): await _register_session(session_registry, available_apis={'run_append_result': True}) await _create_run(db_engine) handler = _handler(db_engine) run_append_result = handler.actions[PluginToRuntimeAction.RUN_APPEND_RESULT.value] + before = await ledger_snapshot() result = await run_append_result( { @@ -1089,14 +1111,16 @@ async def test_run_append_result_rejects_unknown_result_type(session_registry, d assert result.code != 0 assert 'unknown result type' in result.message + assert await ledger_snapshot() == before @pytest.mark.asyncio -async def test_run_append_result_validates_known_payloads(session_registry, db_engine): +async def test_run_append_result_validates_known_payloads(session_registry, db_engine, ledger_snapshot): await _register_session(session_registry, available_apis={'run_append_result': True}) await _create_run(db_engine) handler = _handler(db_engine) run_append_result = handler.actions[PluginToRuntimeAction.RUN_APPEND_RESULT.value] + before = await ledger_snapshot() result = await run_append_result( { @@ -1112,10 +1136,13 @@ async def test_run_append_result_validates_known_payloads(session_registry, db_e assert result.code != 0 assert 'invalid action.requested payload' in result.message + assert await ledger_snapshot() == before @pytest.mark.asyncio -async def test_run_append_and_finalize_claimed_target_require_active_claim(session_registry, db_engine): +async def test_run_append_and_finalize_claimed_target_require_active_claim( + session_registry, db_engine, ledger_snapshot +): await _register_session( session_registry, available_apis={'run_append_result': True, 'run_finalize': True}, @@ -1139,6 +1166,7 @@ async def test_run_append_and_finalize_claimed_target_require_active_claim(sessi handler = _handler(db_engine) run_append_result = handler.actions[PluginToRuntimeAction.RUN_APPEND_RESULT.value] run_finalize = handler.actions[PluginToRuntimeAction.RUN_FINALIZE.value] + before = await ledger_snapshot() missing_claim = await run_append_result( { @@ -1152,6 +1180,9 @@ async def test_run_append_and_finalize_claimed_target_require_active_claim(sessi }, } ) + assert missing_claim.code != 0 + assert 'active claim ownership' in missing_claim.message + assert await ledger_snapshot() == before wrong_claim = await run_finalize( { 'run_id': 'run_1', @@ -1162,6 +1193,9 @@ async def test_run_append_and_finalize_claimed_target_require_active_claim(sessi 'status': 'completed', } ) + assert wrong_claim.code != 0 + assert 'claim ownership is not active' in wrong_claim.message + assert await ledger_snapshot() == before appended = await run_append_result( { 'run_id': 'run_1', @@ -1187,10 +1221,6 @@ async def test_run_append_and_finalize_claimed_target_require_active_claim(sessi } ) - assert missing_claim.code != 0 - assert 'active claim ownership' in missing_claim.message - assert wrong_claim.code != 0 - assert 'claim ownership is not active' in wrong_claim.message assert appended.code == 0 assert finalized.code == 0 assert finalized.data['status'] == 'completed' @@ -1425,7 +1455,7 @@ async def test_non_admin_run_claim_is_scoped_to_session_runner_and_conversation( @pytest.mark.asyncio -async def test_non_admin_renew_and_release_reject_cross_runner_claim(session_registry, db_engine): +async def test_non_admin_renew_and_release_reject_cross_runner_claim(session_registry, db_engine, ledger_snapshot): await _register_session( session_registry, available_apis={'run_renew_claim': True, 'run_release_claim': True}, @@ -1447,6 +1477,7 @@ async def test_non_admin_renew_and_release_reject_cross_runner_claim(session_reg handler = _handler(db_engine) run_renew_claim = handler.actions[PluginToRuntimeAction.RUN_RENEW_CLAIM.value] run_release_claim = handler.actions[PluginToRuntimeAction.RUN_RELEASE_CLAIM.value] + before = await ledger_snapshot() renewed = await run_renew_claim( { @@ -1457,6 +1488,9 @@ async def test_non_admin_renew_and_release_reject_cross_runner_claim(session_reg 'claim_token': claim['claim_token'], } ) + assert renewed.code != 0 + assert 'not accessible' in renewed.message.lower() + assert await ledger_snapshot() == before released = await run_release_claim( { 'run_id': 'run_1', @@ -1467,14 +1501,13 @@ async def test_non_admin_renew_and_release_reject_cross_runner_claim(session_reg } ) - assert renewed.code != 0 - assert 'not accessible' in renewed.message.lower() assert released.code != 0 assert 'not accessible' in released.message.lower() + assert await ledger_snapshot() == before @pytest.mark.asyncio -async def test_non_admin_release_claim_cannot_finalize_run(session_registry, db_engine): +async def test_non_admin_release_claim_cannot_finalize_run(session_registry, db_engine, ledger_snapshot): await _register_session( session_registry, available_apis={'run_claim': True, 'run_release_claim': True}, @@ -1503,6 +1536,7 @@ async def test_non_admin_release_claim_cannot_finalize_run(session_registry, db_ } ) assert claimed.code == 0 + before = await ledger_snapshot() released = await run_release_claim( { @@ -1517,3 +1551,4 @@ async def test_non_admin_release_claim_cannot_finalize_run(session_registry, db_ assert released.code != 0 assert 'use run_finalize' in released.message + assert await ledger_snapshot() == before diff --git a/tests/unit_tests/agent/test_state_store.py b/tests/unit_tests/agent/test_state_store.py index a92c6a4af..8f9e758f1 100644 --- a/tests/unit_tests/agent/test_state_store.py +++ b/tests/unit_tests/agent/test_state_store.py @@ -5,6 +5,7 @@ from __future__ import annotations import asyncio import os import tempfile +from copy import deepcopy import pytest from sqlalchemy.ext.asyncio import create_async_engine @@ -132,6 +133,24 @@ class TestStateScopeHelpers: assert keys['subject'].startswith('subject:v2:') assert keys['runner'].startswith('runner:v2:') assert len(set(keys.values())) == len(keys) + for dimension, value in [ + ('workspace_id', 'ws_other'), + ('bot_id', 'bot_other'), + ]: + other = deepcopy(event) + setattr(other, dimension, value) + for scope in VALID_STATE_SCOPES: + assert build_state_scope_key(scope, other, binding, descriptor) != keys[scope] + for scope, dimension, value in [ + ('conversation', 'conversation_id', 'conv_other'), + ('conversation', 'thread_id', 'thread_other'), + ('actor', 'actor_id', 'user_other'), + ('subject', 'subject_id', 'msg_other'), + ]: + other = deepcopy(event) + owner = other.actor if scope == 'actor' else other.subject if scope == 'subject' else other + setattr(owner, dimension, value) + assert build_state_scope_key(scope, other, binding, descriptor) != keys[scope] def test_scope_key_missing_identity_returns_none(self): descriptor = make_descriptor() diff --git a/tests/unit_tests/plugin/conftest.py b/tests/unit_tests/plugin/conftest.py new file mode 100644 index 000000000..3285bb2a0 --- /dev/null +++ b/tests/unit_tests/plugin/conftest.py @@ -0,0 +1,38 @@ +"""Database fixture for plugin persistence behavior tests.""" + +import pytest +from sqlalchemy import text +from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine + +from langbot.pkg.entity.persistence.base import Base +from langbot.pkg.entity.persistence.bstorage import BinaryStorage +from langbot.pkg.entity.persistence.plugin import PluginSetting +from langbot.pkg.entity.persistence.user import User +from langbot.pkg.entity.persistence.workspace import Workspace + + +@pytest.fixture +async def plugin_db(): + engine = create_async_engine('sqlite+aiosqlite:///:memory:') + try: + async with engine.begin() as connection: + await connection.execute(text('PRAGMA foreign_keys=ON')) + await connection.run_sync( + Base.metadata.create_all, + tables=[User.__table__, Workspace.__table__, PluginSetting.__table__, BinaryStorage.__table__], + ) + async with AsyncSession(engine, expire_on_commit=False) as session: + session.add_all( + Workspace( + uuid=workspace, + instance_uuid='instance-a', + name=workspace, + slug=workspace, + source='cloud_projection', + ) + for workspace in ('workspace-a', 'workspace-b') + ) + await session.flush() + yield session + finally: + await engine.dispose() diff --git a/tests/unit_tests/plugin/test_connector_pure.py b/tests/unit_tests/plugin/test_connector_pure.py index 4ed20f1cc..5f671fc8d 100644 --- a/tests/unit_tests/plugin/test_connector_pure.py +++ b/tests/unit_tests/plugin/test_connector_pure.py @@ -123,6 +123,7 @@ class TestExtractDepsMetadata: # Should find requirements.txt in subdirectory assert task_context.metadata['deps_total'] == 2 + assert task_context.metadata['deps_list'] == ['pytest', 'black'] def test_archive_preview_rejects_extreme_compression_ratio(self): from langbot.pkg.plugin.connector import inspect_plugin_archive_metadata diff --git a/tests/unit_tests/plugin/test_connector_reconcile.py b/tests/unit_tests/plugin/test_connector_reconcile.py index 0b54f9384..bfb7c2431 100644 --- a/tests/unit_tests/plugin/test_connector_reconcile.py +++ b/tests/unit_tests/plugin/test_connector_reconcile.py @@ -583,11 +583,15 @@ async def test_local_install_cleans_untracked_legacy_plugin_before_runtime_apply @pytest.mark.asyncio -@pytest.mark.parametrize(('remaining_references', 'statement_count'), [(1, 1), (0, 2)]) +@pytest.mark.parametrize('remaining_references', [1, 0]) async def test_artifact_cleanup_is_reference_counted_within_workspace( remaining_references: int, - statement_count: int, + plugin_db, ): + import sqlalchemy + from langbot.pkg.entity.persistence.plugin import PluginSetting + from langbot.pkg.entity.persistence.bstorage import BinaryStorage + app = SimpleNamespace( instance_config=SimpleNamespace(data={'plugin': {'enable': True}}), deployment=SimpleNamespace(mode='cloud'), @@ -598,19 +602,38 @@ async def test_artifact_cleanup_is_reference_counted_within_workspace( workspace_uuid='workspace-a', placement_generation=1, ) - statements = [] - - async def execute(statement): - statements.append(statement) - return SimpleNamespace(scalar_one=lambda: remaining_references) + digest = 'a' * 64 + settings = [('workspace-b', digest), ('workspace-a', 'b' * 64)] + if remaining_references: + settings.append(('workspace-a', digest)) + plugin_db.add_all( + PluginSetting(workspace_uuid=workspace, plugin_author='author', plugin_name=artifact, artifact_digest=artifact) + for workspace, artifact in settings + ) + plugin_db.add_all( + BinaryStorage( + workspace_uuid=workspace, unique_key=key, key=key, owner_type=owner_type, owner=owner, value=b'pkg' + ) + for workspace, key, owner_type, owner in [ + ('workspace-a', 'target', 'plugin_artifact', digest), + ('workspace-b', 'other-workspace', 'plugin_artifact', digest), + ('workspace-a', 'other-artifact', 'plugin_artifact', 'b' * 64), + ('workspace-a', 'other-owner-type', 'plugin', digest), + ] + ) + await plugin_db.flush() await connector._delete_artifact_if_unreferenced( execution_context, - 'a' * 64, - execute=execute, + digest, + execute=plugin_db.execute, ) - assert len(statements) == statement_count + remaining = (await plugin_db.execute(sqlalchemy.select(BinaryStorage.unique_key))).scalars().all() + expected = {'other-workspace', 'other-artifact', 'other-owner-type'} + if remaining_references: + expected.add('target') + assert set(remaining) == expected @pytest.mark.asyncio diff --git a/tests/unit_tests/plugin/test_extract_deps.py b/tests/unit_tests/plugin/test_extract_deps.py index 0980f4cc3..0753e6b3e 100644 --- a/tests/unit_tests/plugin/test_extract_deps.py +++ b/tests/unit_tests/plugin/test_extract_deps.py @@ -113,6 +113,7 @@ numpy""" # Should find nested requirements.txt (ends with 'requirements.txt') assert task_context.metadata.get('deps_total') == 2 + assert task_context.metadata['deps_list'] == ['requests', 'flask'] def test_no_requirements_in_zip(self): """Test handling ZIP without requirements.txt.""" diff --git a/tests/unit_tests/plugin/test_handler_actions.py b/tests/unit_tests/plugin/test_handler_actions.py index 7f0c10ee5..8fd068745 100644 --- a/tests/unit_tests/plugin/test_handler_actions.py +++ b/tests/unit_tests/plugin/test_handler_actions.py @@ -236,10 +236,16 @@ class TestInitializePluginSettings: app.persistence_mgr.execute_async.assert_awaited_once() @pytest.mark.asyncio - async def test_existing_desired_setting_remains_core_owned(self, app): + async def test_existing_desired_setting_remains_core_owned(self, app, plugin_db): """Runtime initialization validates identity without rewriting Core state.""" + import sqlalchemy + from langbot.pkg.entity.persistence.plugin import PluginSetting + runtime_handler = make_handler(app) - existing_setting = SimpleNamespace( + existing_setting = PluginSetting( + workspace_uuid='workspace-a', + plugin_author='test-author', + plugin_name='test-plugin', enabled=False, priority=5, config={'key': 'value'}, @@ -247,7 +253,29 @@ class TestInitializePluginSettings: runtime_revision=1, artifact_digest='a' * 64, ) - app.persistence_mgr.execute_async.return_value = make_result(existing_setting) + other_setting = PluginSetting( + workspace_uuid='workspace-b', + plugin_author='test-author', + plugin_name='test-plugin', + enabled=True, + ) + plugin_db.add_all([other_setting, existing_setting]) + await plugin_db.flush() + + async def execute(statement): + result = await plugin_db.execute(statement) + return result.scalars() if statement.is_select or statement.is_update else result + + app.persistence_mgr.execute_async.side_effect = execute + before = ( + ( + await plugin_db.execute( + sqlalchemy.select(*PluginSetting.__table__.c).order_by(PluginSetting.workspace_uuid) + ) + ) + .mappings() + .all() + ) response = await runtime_handler.actions[RuntimeToLangBotAction.INITIALIZE_PLUGIN_SETTINGS.value]( { @@ -260,6 +288,17 @@ class TestInitializePluginSettings: assert response.code == 0 app.persistence_mgr.execute_async.assert_awaited_once() + plugin_db.expire_all() + after = ( + ( + await plugin_db.execute( + sqlalchemy.select(*PluginSetting.__table__.c).order_by(PluginSetting.workspace_uuid) + ) + ) + .mappings() + .all() + ) + assert after == before class TestSetBinaryStorage: diff --git a/tests/unit_tests/plugin/test_handler_helpers.py b/tests/unit_tests/plugin/test_handler_helpers.py index b126a9b18..145d6e5bd 100644 --- a/tests/unit_tests/plugin/test_handler_helpers.py +++ b/tests/unit_tests/plugin/test_handler_helpers.py @@ -107,41 +107,69 @@ class TestMakeRagErrorResponse: class TestCleanupPluginData: """Tests for cleanup_plugin_data method.""" - @pytest.mark.asyncio - async def test_deletes_plugin_settings(self): - """Test that plugin settings are deleted.""" - handler_module = get_handler_module() + @pytest.fixture + async def cleanup(self, plugin_db): + from langbot.pkg.entity.persistence.plugin import PluginSetting + from langbot.pkg.entity.persistence.bstorage import BinaryStorage + from langbot_plugin.entities.io.context import InstallationBinding - mock_app = Mock() - mock_app.persistence_mgr = AsyncMock() - mock_app.persistence_mgr.execute_async = AsyncMock() - - # Mock the handler without connection (we only need ap) - handler_instance = Mock(spec=handler_module.RuntimeConnectionHandler) - handler_instance.ap = mock_app - - # Call cleanup_plugin_data - await handler_module.RuntimeConnectionHandler.cleanup_plugin_data( - handler_instance, 'test-author', 'test-plugin' + plugin_db.add_all( + PluginSetting(workspace_uuid=workspace, plugin_author=author, plugin_name=name) + for workspace, author, name in [ + ('workspace-a', 'test-author', 'test-plugin'), + ('workspace-b', 'test-author', 'test-plugin'), + ('workspace-a', 'other-author', 'test-plugin'), + ('workspace-a', 'test-author', 'other-plugin'), + ] ) - - # Verify plugin settings delete was called - calls = mock_app.persistence_mgr.execute_async.call_args_list - assert len(calls) >= 1 + plugin_db.add_all( + BinaryStorage( + workspace_uuid=workspace, unique_key=key, key=key, owner_type=owner_type, owner=owner, value=b'data' + ) + for workspace, key, owner_type, owner in [ + ('workspace-a', 'target', 'plugin', 'test-author/test-plugin'), + ('workspace-b', 'other-workspace', 'plugin', 'test-author/test-plugin'), + ('workspace-a', 'other-owner', 'plugin', 'other-author/test-plugin'), + ('workspace-a', 'other-type', 'skill', 'test-author/test-plugin'), + ] + ) + await plugin_db.flush() + app = Mock() + app.persistence_mgr.execute_async = plugin_db.execute + handler = get_handler_module().RuntimeConnectionHandler(Mock(), AsyncMock(), app) + binding = InstallationBinding( + instance_uuid='instance-a', + workspace_uuid='workspace-a', + placement_generation=1, + installation_uuid='00000000-0000-4000-8000-000000000001', + runtime_revision=1, + artifact_digest='a' * 64, + ) + with handler.installation_scope(binding): + await handler.cleanup_plugin_data('test-author', 'test-plugin') + return plugin_db @pytest.mark.asyncio - async def test_deletes_binary_storage(self): + async def test_deletes_plugin_settings(self, cleanup): + import sqlalchemy + from langbot.pkg.entity.persistence.plugin import PluginSetting + + rows = ( + await cleanup.execute( + sqlalchemy.select(PluginSetting.workspace_uuid, PluginSetting.plugin_author, PluginSetting.plugin_name) + ) + ).all() + assert set(rows) == { + ('workspace-b', 'test-author', 'test-plugin'), + ('workspace-a', 'other-author', 'test-plugin'), + ('workspace-a', 'test-author', 'other-plugin'), + } + + @pytest.mark.asyncio + async def test_deletes_binary_storage(self, cleanup): """Test that binary storage is deleted.""" - handler_module = get_handler_module() + import sqlalchemy + from langbot.pkg.entity.persistence.bstorage import BinaryStorage - mock_app = Mock() - mock_app.persistence_mgr = AsyncMock() - mock_app.persistence_mgr.execute_async = AsyncMock() - - handler_instance = Mock(spec=handler_module.RuntimeConnectionHandler) - handler_instance.ap = mock_app - - await handler_module.RuntimeConnectionHandler.cleanup_plugin_data(handler_instance, 'author', 'plugin-name') - - # Should have at least 2 calls: one for settings, one for binary storage - assert mock_app.persistence_mgr.execute_async.call_count >= 2 + keys = (await cleanup.execute(sqlalchemy.select(BinaryStorage.unique_key))).scalars().all() + assert set(keys) == {'other-workspace', 'other-owner', 'other-type'}