From 7c39df1d4b0af03aadce07f29a5beb49ea9d319a Mon Sep 17 00:00:00 2001 From: Chuck Butkus Date: Mon, 3 Nov 2025 12:23:46 -0500 Subject: [PATCH] Add back isolation test --- ..._saas_sql_app_conversation_info_service.py | 191 ++++++++++++++++-- 1 file changed, 177 insertions(+), 14 deletions(-) diff --git a/enterprise/tests/unit/storage/test_saas_sql_app_conversation_info_service.py b/enterprise/tests/unit/storage/test_saas_sql_app_conversation_info_service.py index c6df81f630..d768a352ca 100644 --- a/enterprise/tests/unit/storage/test_saas_sql_app_conversation_info_service.py +++ b/enterprise/tests/unit/storage/test_saas_sql_app_conversation_info_service.py @@ -4,17 +4,21 @@ This module tests the SAAS implementation of SQLAppConversationInfoService, focusing on user isolation, SAAS metadata handling, and multi-tenant functionality. """ -from unittest.mock import AsyncMock, MagicMock, patch +from datetime import datetime, timezone +from typing import AsyncGenerator +from unittest.mock import AsyncMock, MagicMock from uuid import UUID, uuid4 import pytest - from openhands.app_server.app_conversation.app_conversation_models import ( AppConversationInfo, ) from openhands.app_server.user.specifiy_user_context import SpecifyUserContext +from openhands.app_server.utils.sql_utils import Base from openhands.integrations.service_types import ProviderType from openhands.storage.data_models.conversation_metadata import ConversationTrigger +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine +from sqlalchemy.pool import StaticPool # Import the SAAS service from enterprise.storage.saas_app_conversation_info_injector import ( @@ -22,6 +26,98 @@ from enterprise.storage.saas_app_conversation_info_injector import ( ) +@pytest.fixture +async def async_engine(): + """Create an async SQLite engine for testing.""" + engine = create_async_engine( + 'sqlite+aiosqlite:///:memory:', + poolclass=StaticPool, + connect_args={'check_same_thread': False}, + echo=False, + ) + + # Create all tables + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + + yield engine + + await engine.dispose() + + +@pytest.fixture +async def async_session(async_engine) -> AsyncGenerator[AsyncSession, None]: + """Create an async session for testing.""" + async_session_maker = async_sessionmaker( + async_engine, class_=AsyncSession, expire_on_commit=False + ) + + async with async_session_maker() as db_session: + yield db_session + + +@pytest.fixture +def service(async_session) -> SaasSQLAppConversationInfoService: + """Create a SQLAppConversationInfoService instance for testing.""" + return SaasSQLAppConversationInfoService( + db_session=async_session, user_context=SpecifyUserContext(user_id=None) + ) + + +@pytest.fixture +def service_with_user(async_session) -> SaasSQLAppConversationInfoService: + """Create a SQLAppConversationInfoService instance with a user_id for testing.""" + return SaasSQLAppConversationInfoService( + db_session=async_session, + user_context=SpecifyUserContext(user_id='11111111-1111-1111-1111-111111111111'), + ) + + +@pytest.fixture +def sample_conversation_info() -> AppConversationInfo: + """Create a sample AppConversationInfo for testing.""" + return AppConversationInfo( + id=uuid4(), + created_by_user_id='11111111-1111-1111-1111-111111111111', + sandbox_id='sandbox_123', + selected_repository='https://github.com/test/repo', + selected_branch='main', + git_provider=ProviderType.GITHUB, + title='Test Conversation', + trigger=ConversationTrigger.GUI, + pr_number=[123, 456], + llm_model='gpt-4', + metrics=None, + created_at=datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc), + updated_at=datetime(2024, 1, 1, 12, 30, 0, tzinfo=timezone.utc), + ) + + +@pytest.fixture +def multiple_conversation_infos() -> list[AppConversationInfo]: + """Create multiple AppConversationInfo instances for testing.""" + base_time = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc) + + return [ + AppConversationInfo( + id=uuid4(), + created_by_user_id=None, + sandbox_id=f'sandbox_{i}', + selected_repository=f'https://github.com/test/repo{i}', + selected_branch='main', + git_provider=ProviderType.GITHUB, + title=f'Test Conversation {i}', + trigger=ConversationTrigger.GUI, + pr_number=[i * 100], + llm_model='gpt-4', + metrics=None, + created_at=base_time.replace(hour=12 + i), + updated_at=base_time.replace(hour=12 + i, minute=30), + ) + for i in range(1, 6) # Create 5 conversations + ] + + @pytest.fixture def mock_db_session(): """Create a mock database session.""" @@ -44,8 +140,7 @@ def user2_context(): def saas_service_user1(mock_db_session, user1_context): """Create a SaasSQLAppConversationInfoService instance for user1.""" return SaasSQLAppConversationInfoService( - db_session=mock_db_session, - user_context=user1_context + db_session=mock_db_session, user_context=user1_context ) @@ -53,8 +148,7 @@ def saas_service_user1(mock_db_session, user1_context): def saas_service_user2(mock_db_session, user2_context): """Create a SaasSQLAppConversationInfoService instance for user2.""" return SaasSQLAppConversationInfoService( - db_session=mock_db_session, - user_context=user2_context + db_session=mock_db_session, user_context=user2_context ) @@ -79,7 +173,7 @@ class TestSaasSQLAppConversationInfoService: """Test that different service instances have different user contexts.""" user1_id = await saas_service_user1.user_context.get_user_id() user2_id = await saas_service_user2.user_context.get_user_id() - + assert user1_id == '11111111-1111-1111-1111-111111111111' assert user2_id == '22222222-2222-2222-2222-222222222222' assert user1_id != user2_id @@ -102,8 +196,10 @@ class TestSaasSQLAppConversationInfoService: ): """Test that _to_info_with_user_id properly sets user_id from SAAS metadata.""" from storage.stored_conversation_metadata import StoredConversationMetadata - from storage.stored_conversation_metadata_saas import StoredConversationMetadataSaas - + from storage.stored_conversation_metadata_saas import ( + StoredConversationMetadataSaas, + ) + # Create mock metadata objects stored_metadata = MagicMock(spec=StoredConversationMetadata) stored_metadata.conversation_id = '12345678-1234-5678-1234-567812345678' @@ -116,6 +212,7 @@ class TestSaasSQLAppConversationInfoService: stored_metadata.pr_number = [] stored_metadata.llm_model = None from datetime import datetime, timezone + stored_metadata.created_at = datetime.now(timezone.utc) stored_metadata.last_updated_at = datetime.now(timezone.utc) stored_metadata.accumulated_cost = 0.0 @@ -128,14 +225,80 @@ class TestSaasSQLAppConversationInfoService: stored_metadata.reasoning_tokens = 0 stored_metadata.context_window = 0 stored_metadata.per_turn_token = 0 - + saas_metadata = MagicMock(spec=StoredConversationMetadataSaas) saas_metadata.user_id = UUID('11111111-1111-1111-1111-111111111111') - + # Test the _to_info_with_user_id method - result = saas_service_user1._to_info_with_user_id(stored_metadata, saas_metadata) - + result = saas_service_user1._to_info_with_user_id( + stored_metadata, saas_metadata + ) + # Verify that the user_id from SAAS metadata is used assert result.created_by_user_id == '11111111-1111-1111-1111-111111111111' assert result.title == 'Test Conversation' - assert result.sandbox_id == 'test-sandbox' \ No newline at end of file + assert result.sandbox_id == 'test-sandbox' + + @pytest.mark.asyncio + async def test_user_isolation( + self, + async_session: AsyncSession, + multiple_conversation_infos: list[AppConversationInfo], + ): + """Test that user isolation works correctly.""" + # Create services for different users + user1_service = SaasSQLAppConversationInfoService( + db_session=async_session, + user_context=SpecifyUserContext( + user_id='11111111-1111-1111-1111-111111111111' + ), + ) + user2_service = SaasSQLAppConversationInfoService( + db_session=async_session, + user_context=SpecifyUserContext( + user_id='22222222-2222-2222-2222-222222222222' + ), + ) + + # Create conversations for different users + user1_info = AppConversationInfo( + id=uuid4(), + created_by_user_id='11111111-1111-1111-1111-111111111111', + sandbox_id='sandbox_user1', + title='User 1 Conversation', + ) + + user2_info = AppConversationInfo( + id=uuid4(), + created_by_user_id='22222222-2222-2222-2222-222222222222', + sandbox_id='sandbox_user2', + title='User 2 Conversation', + ) + + # Save conversations + await user1_service.save_app_conversation_info(user1_info) + await user2_service.save_app_conversation_info(user2_info) + + # User 1 should only see their conversation + user1_page = await user1_service.search_app_conversation_info() + assert len(user1_page.items) == 1 + assert ( + user1_page.items[0].created_by_user_id + == '11111111-1111-1111-1111-111111111111' + ) + + # User 2 should only see their conversation + user2_page = await user2_service.search_app_conversation_info() + assert len(user2_page.items) == 1 + assert ( + user2_page.items[0].created_by_user_id + == '22222222-2222-2222-2222-222222222222' + ) + + # User 1 should not be able to get user 2's conversation + user2_from_user1 = await user1_service.get_app_conversation_info(user2_info.id) + assert user2_from_user1 is None + + # User 2 should not be able to get user 1's conversation + user1_from_user2 = await user2_service.get_app_conversation_info(user1_info.id) + assert user1_from_user2 is None