import uuid from contextlib import asynccontextmanager from unittest.mock import AsyncMock, MagicMock, patch import pytest from server.routes.org_models import OrgUpdate from sqlalchemy import select from sqlalchemy.exc import IntegrityError from storage.org import Org from storage.org_invitation import OrgInvitation from storage.org_member import OrgMember from storage.org_store import OrgStore from storage.role import Role from storage.user import User from openhands.app_server.settings.settings_models import Settings from openhands.sdk.settings import ConversationSettings, OpenHandsAgentSettings @pytest.fixture def mock_litellm_api(): api_key_patch = patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test_key') api_url_patch = patch( 'storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.url' ) team_id_patch = patch('storage.lite_llm_manager.LITE_LLM_TEAM_ID', 'test_team') client_patch = patch('httpx.AsyncClient') with api_key_patch, api_url_patch, team_id_patch, client_patch as mock_client: mock_response = AsyncMock() mock_response.is_success = True mock_response.json = MagicMock(return_value={'key': 'test_api_key'}) mock_client.return_value.__aenter__.return_value.post.return_value = ( mock_response ) mock_client.return_value.__aenter__.return_value.get.return_value = ( mock_response ) mock_client.return_value.__aenter__.return_value.patch.return_value = ( mock_response ) yield mock_client @pytest.mark.asyncio async def test_get_org_by_id(async_session_maker, mock_litellm_api): # Test getting org by ID async with async_session_maker() as session: # Create a test org org = Org(name='test-org') session.add(org) await session.commit() await session.refresh(org) org_id = org.id # Test retrieval with ( patch('storage.org_store.a_session_maker', async_session_maker), ): retrieved_org = await OrgStore.get_org_by_id(org_id) assert retrieved_org is not None assert retrieved_org.id == org_id assert retrieved_org.name == 'test-org' @pytest.mark.asyncio async def test_get_org_by_id_not_found(async_session_maker): # Test getting org by ID when it doesn't exist with patch('storage.org_store.a_session_maker', async_session_maker): non_existent_id = uuid.uuid4() retrieved_org = await OrgStore.get_org_by_id(non_existent_id) assert retrieved_org is None @pytest.mark.asyncio async def test_list_orgs(async_session_maker, mock_litellm_api): # Test listing all orgs async with async_session_maker() as session: # Create test orgs org1 = Org(name='test-org-1') org2 = Org(name='test-org-2') session.add_all([org1, org2]) await session.commit() # Test listing with ( patch('storage.org_store.a_session_maker', async_session_maker), ): orgs = await OrgStore.list_orgs() assert len(orgs) >= 2 org_names = [org.name for org in orgs] assert 'test-org-1' in org_names assert 'test-org-2' in org_names @pytest.mark.asyncio async def test_update_org(async_session_maker, mock_litellm_api): # Test updating org details async with async_session_maker() as session: # Create a test org org = Org( name='test-org', agent_settings=OpenHandsAgentSettings(agent='CodeActAgent'), ) session.add(org) await session.commit() await session.refresh(org) org_id = org.id # Test update with ( patch('storage.org_store.a_session_maker', async_session_maker), patch( 'storage.org_store.OrgStore._maybe_get_managed_llm_key_for_user', new=AsyncMock(return_value=None), ), patch( 'storage.org_member_store.OrgMemberStore.update_all_members_settings_async', new=AsyncMock(), ), ): updated_org = await OrgStore.update_org( org_id=org_id, update_data=OrgUpdate( name='updated-org', agent_settings_diff={'llm': {'model': 'openhands/claude-3'}}, ), user_id=str(uuid.uuid4()), ) assert updated_org is not None assert updated_org.name == 'updated-org' agent_settings = OrgStore.get_agent_settings_from_org(updated_org) assert agent_settings.llm.model == 'litellm_proxy/claude-3' def test_get_org_settings_from_org_use_persisted_loaders(): org = MagicMock(spec=Org) org.agent_settings = {'legacy': True} org.conversation_settings = {'legacy': True} loaded_agent_settings = OpenHandsAgentSettings(agent='MigratedAgent') loaded_conversation_settings = ConversationSettings(max_iterations=77) with ( patch( 'storage.org_store._load_persisted_agent_settings', return_value=loaded_agent_settings, ) as agent_loader, patch( 'storage.org_store._load_persisted_conversation_settings', return_value=loaded_conversation_settings, ) as conversation_loader, ): assert OrgStore.get_agent_settings_from_org(org).agent == 'MigratedAgent' assert OrgStore.get_conversation_settings_from_org(org).max_iterations == 77 agent_loader.assert_called_once_with({'legacy': True}) conversation_loader.assert_called_once_with({'legacy': True}) @pytest.mark.asyncio async def test_update_org_not_found(async_session_maker): # Test updating org that doesn't exist with patch('storage.org_store.a_session_maker', async_session_maker): from uuid import uuid4 updated_org = await OrgStore.update_org( org_id=uuid4(), update_data=OrgUpdate(name='updated-org') ) assert updated_org is None @pytest.mark.asyncio async def test_create_org(async_session_maker, mock_litellm_api): # Test creating a new org with ( patch('storage.org_store.a_session_maker', async_session_maker), ): org = await OrgStore.create_org( kwargs={ 'name': 'new-org', 'agent_settings': OpenHandsAgentSettings(agent='CodeActAgent'), } ) assert org is not None assert org.name == 'new-org' assert org.agent_settings['agent'] == 'CodeActAgent' assert org.id is not None @pytest.mark.asyncio async def test_create_org_v1_enabled_defaults_to_true_when_default_is_true( async_session_maker, mock_litellm_api ): """ GIVEN: DEFAULT_V1_ENABLED is True and org.v1_enabled is not specified (None) WHEN: create_org is called THEN: org.v1_enabled should be set to True """ with ( patch('storage.org_store.a_session_maker', async_session_maker), patch('storage.org_store.DEFAULT_V1_ENABLED', True), ): org = await OrgStore.create_org(kwargs={'name': 'test-org-v1-default-true'}) assert org is not None assert org.v1_enabled is True @pytest.mark.asyncio async def test_create_org_v1_enabled_defaults_to_false_when_default_is_false( async_session_maker, mock_litellm_api ): """ GIVEN: DEFAULT_V1_ENABLED is False and org.v1_enabled is not specified (None) WHEN: create_org is called THEN: org.v1_enabled should be set to False """ with ( patch('storage.org_store.a_session_maker', async_session_maker), patch('storage.org_store.DEFAULT_V1_ENABLED', False), ): org = await OrgStore.create_org(kwargs={'name': 'test-org-v1-default-false'}) assert org is not None assert org.v1_enabled is False @pytest.mark.asyncio async def test_create_org_v1_enabled_explicit_false_overrides_default_true( async_session_maker, mock_litellm_api ): """ GIVEN: DEFAULT_V1_ENABLED is True but org.v1_enabled is explicitly set to False WHEN: create_org is called THEN: org.v1_enabled should stay False (explicit value wins over default) """ with ( patch('storage.org_store.a_session_maker', async_session_maker), patch('storage.org_store.DEFAULT_V1_ENABLED', True), ): org = await OrgStore.create_org( kwargs={'name': 'test-org-v1-explicit-false', 'v1_enabled': False} ) assert org is not None assert org.v1_enabled is False @pytest.mark.asyncio async def test_create_org_v1_enabled_explicit_true_overrides_default_false( async_session_maker, mock_litellm_api ): """ GIVEN: DEFAULT_V1_ENABLED is False but org.v1_enabled is explicitly set to True WHEN: create_org is called THEN: org.v1_enabled should stay True (explicit value wins over default) """ with ( patch('storage.org_store.a_session_maker', async_session_maker), patch('storage.org_store.DEFAULT_V1_ENABLED', False), ): org = await OrgStore.create_org( kwargs={'name': 'test-org-v1-explicit-true', 'v1_enabled': True} ) assert org is not None assert org.v1_enabled is True @pytest.mark.asyncio async def test_get_org_by_name(async_session_maker, mock_litellm_api): # Test getting org by name async with async_session_maker() as session: # Create a test org org = Org(name='test-org-by-name') session.add(org) await session.commit() # Test retrieval with ( patch('storage.org_store.a_session_maker', async_session_maker), ): retrieved_org = await OrgStore.get_org_by_name('test-org-by-name') assert retrieved_org is not None assert retrieved_org.name == 'test-org-by-name' @pytest.mark.asyncio async def test_get_current_org_from_keycloak_user_id( async_session_maker, mock_litellm_api ): # Test getting current org from user ID test_user_id = uuid.uuid4() async with async_session_maker() as session: # Create test data org = Org(name='test-org') session.add(org) await session.flush() from storage.user import User user = User(id=test_user_id, current_org_id=org.id) session.add(user) await session.commit() await session.refresh(org) # Test retrieval with ( patch('storage.org_store.a_session_maker', async_session_maker), ): retrieved_org = await OrgStore.get_current_org_from_keycloak_user_id( str(test_user_id) ) assert retrieved_org is not None assert retrieved_org.name == 'test-org' def test_get_kwargs_from_settings(): # Test extracting org kwargs from settings settings = Settings() settings.update( { 'language': 'es', 'enable_sound_notifications': True, 'agent_settings_diff': { 'agent': 'CodeActAgent', 'llm': { 'model': 'anthropic/claude-sonnet-4-5-20250929', 'api_key': 'test-key', }, }, } ) kwargs = OrgStore.get_kwargs_from_settings(settings) # Should only include fields that exist in Org model assert 'agent_settings' in kwargs assert 'agent' not in kwargs assert 'default_llm_model' not in kwargs assert kwargs['agent_settings']['agent'] == 'CodeActAgent' assert ( kwargs['agent_settings']['llm']['model'] == 'anthropic/claude-sonnet-4-5-20250929' ) # Should not include fields that don't exist in Org model assert 'language' not in kwargs # language is not in Org model assert 'llm_api_key' not in kwargs assert 'llm_model' not in kwargs assert 'enable_sound_notifications' not in kwargs @pytest.mark.asyncio async def test_persist_org_with_owner_success(async_session_maker, mock_litellm_api): """ GIVEN: Valid org and org_member entities WHEN: persist_org_with_owner is called THEN: Both entities are persisted in a single transaction and org is returned """ # Arrange org_id = uuid.uuid4() user_id = uuid.uuid4() # Create user and role first async with async_session_maker() as session: user = User(id=user_id, current_org_id=org_id) role = Role(id=1, name='owner', rank=1) session.add(user) session.add(role) await session.commit() org = Org( id=org_id, name='Test Organization', contact_name='John Doe', contact_email='john@example.com', ) org_member = OrgMember( org_id=org_id, user_id=user_id, role_id=1, status='active', llm_api_key='test-api-key-123', ) # Act with patch('storage.org_store.a_session_maker', async_session_maker): result = await OrgStore.persist_org_with_owner(org, org_member) # Assert assert result is not None assert result.id == org_id assert result.name == 'Test Organization' # Verify both entities were persisted async with async_session_maker() as session: persisted_org = await session.get(Org, org_id) assert persisted_org is not None assert persisted_org.name == 'Test Organization' result = await session.execute( select(OrgMember).filter_by(org_id=org_id, user_id=user_id) ) persisted_member = result.scalars().first() assert persisted_member is not None assert persisted_member.status == 'active' assert persisted_member.role_id == 1 @pytest.mark.asyncio async def test_persist_org_with_owner_returns_refreshed_org( async_session_maker, mock_litellm_api ): """ GIVEN: Valid org and org_member entities WHEN: persist_org_with_owner is called THEN: The returned org is refreshed from database with all fields populated """ # Arrange org_id = uuid.uuid4() user_id = uuid.uuid4() async with async_session_maker() as session: user = User(id=user_id, current_org_id=org_id) role = Role(id=1, name='owner', rank=1) session.add(user) session.add(role) await session.commit() org = Org( id=org_id, name='Test Org', contact_name='Jane Doe', contact_email='jane@example.com', agent_settings=OpenHandsAgentSettings(agent='CodeActAgent'), ) org_member = OrgMember( org_id=org_id, user_id=user_id, role_id=1, status='active', llm_api_key='test-key', ) # Act with patch('storage.org_store.a_session_maker', async_session_maker): result = await OrgStore.persist_org_with_owner(org, org_member) # Assert - verify the returned object has database-generated fields assert result.id == org_id assert result.name == 'Test Org' assert result.agent_settings['agent'] == 'CodeActAgent' # Verify org_version was set by create_org logic (if applicable) assert hasattr(result, 'org_version') @pytest.mark.asyncio async def test_persist_org_with_owner_transaction_atomicity( async_session_maker, mock_litellm_api ): """ GIVEN: Valid org but invalid org_member (missing required field) WHEN: persist_org_with_owner is called THEN: Transaction fails and neither entity is persisted """ # Arrange org_id = uuid.uuid4() user_id = uuid.uuid4() async with async_session_maker() as session: user = User(id=user_id, current_org_id=org_id) role = Role(id=1, name='owner', rank=1) session.add(user) session.add(role) await session.commit() org = Org( id=org_id, name='Test Org', contact_name='John Doe', contact_email='john@example.com', ) # Create invalid org_member (missing required llm_api_key field) org_member = OrgMember( org_id=org_id, user_id=user_id, role_id=1, status='active', # llm_api_key is missing - should cause NOT NULL constraint violation ) # Act & Assert with patch('storage.org_store.a_session_maker', async_session_maker): with pytest.raises(IntegrityError): # NOT NULL constraint violation await OrgStore.persist_org_with_owner(org, org_member) # Verify neither entity was persisted (transaction rolled back) async with async_session_maker() as session: persisted_org = await session.get(Org, org_id) assert persisted_org is None result = await session.execute( select(OrgMember).filter_by(org_id=org_id, user_id=user_id) ) persisted_member = result.scalars().first() assert persisted_member is None @pytest.mark.asyncio async def test_persist_org_with_owner_with_multiple_fields( async_session_maker, mock_litellm_api ): """ GIVEN: Org with multiple optional fields populated WHEN: persist_org_with_owner is called THEN: All fields are persisted correctly """ # Arrange org_id = uuid.uuid4() user_id = uuid.uuid4() async with async_session_maker() as session: user = User(id=user_id, current_org_id=org_id) role = Role(id=1, name='owner', rank=1) session.add(user) session.add(role) await session.commit() org = Org( id=org_id, name='Complex Org', contact_name='Alice Smith', contact_email='alice@example.com', agent_settings=OpenHandsAgentSettings(agent='CodeActAgent'), billing_margin=0.15, ) org_member = OrgMember( org_id=org_id, user_id=user_id, role_id=1, status='active', llm_api_key='test-key', agent_settings_diff={ 'llm': {'model': 'gpt-4'}, }, conversation_settings_diff={ 'max_iterations': 100, }, ) # Act with patch('storage.org_store.a_session_maker', async_session_maker): result = await OrgStore.persist_org_with_owner(org, org_member) # Assert assert result.name == 'Complex Org' assert result.agent_settings['agent'] == 'CodeActAgent' assert result.billing_margin == 0.15 # Verify persistence async with async_session_maker() as session: persisted_org = await session.get(Org, org_id) assert persisted_org.agent_settings['agent'] == 'CodeActAgent' assert persisted_org.billing_margin == 0.15 result_query = await session.execute( select(OrgMember).filter_by(org_id=org_id, user_id=user_id) ) persisted_member = result_query.scalars().first() assert persisted_member.conversation_settings_diff['max_iterations'] == 100 assert persisted_member.agent_settings_diff['llm']['model'] == 'gpt-4' @pytest.mark.asyncio @pytest.mark.skip( reason='Uses PostgreSQL-specific ::uuid cast syntax not supported by SQLite' ) async def test_delete_org_cascade_success(async_session_maker, mock_litellm_api): """ GIVEN: Valid organization with associated data WHEN: delete_org_cascade is called THEN: Organization and all associated data are deleted and org object is returned """ # Arrange org_id = uuid.uuid4() # Create expected return object expected_org = Org( id=org_id, name='Test Organization', contact_name='John Doe', contact_email='john@example.com', ) async with async_session_maker() as session: session.add(expected_org) await session.commit() with patch('storage.org_store.a_session_maker', async_session_maker): # Act result = await OrgStore.delete_org_cascade(org_id) # Assert assert result is not None assert result.id == org_id assert result.name == 'Test Organization' assert result.contact_name == 'John Doe' assert result.contact_email == 'john@example.com' @pytest.mark.asyncio async def test_delete_org_cascade_not_found(async_session_maker): """ GIVEN: Organization ID that doesn't exist WHEN: delete_org_cascade is called THEN: None is returned """ # Arrange non_existent_id = uuid.uuid4() with patch('storage.org_store.a_session_maker', async_session_maker): # Act result = await OrgStore.delete_org_cascade(non_existent_id) # Assert assert result is None @pytest.mark.asyncio async def test_delete_org_cascade_litellm_failure_causes_rollback( async_session_maker, mock_litellm_api ): """ GIVEN: Organization exists but LiteLLM cleanup fails WHEN: delete_org_cascade is called THEN: Transaction is rolled back and organization still exists """ # Arrange org_id = uuid.uuid4() user_id = uuid.uuid4() async with async_session_maker() as session: role = Role(id=1, name='owner', rank=1) user = User(id=user_id, current_org_id=org_id) org = Org( id=org_id, name='Test Organization', contact_name='John Doe', contact_email='john@example.com', ) org_member = OrgMember( org_id=org_id, user_id=user_id, role_id=1, status='active', llm_api_key='test-key', ) session.add_all([role, user, org, org_member]) await session.commit() # Mock delete_org_cascade to simulate LiteLLM failure litellm_error = Exception('LiteLLM API unavailable') async def mock_delete_org_cascade_with_failure(org_id_param): # Verify org exists but then fail with LiteLLM error async with async_session_maker() as session: org = await session.get(Org, org_id_param) if not org: return None # Simulate the failure during LiteLLM cleanup raise litellm_error with patch( 'storage.org_store.OrgStore.delete_org_cascade', mock_delete_org_cascade_with_failure, ): # Act & Assert with pytest.raises(Exception) as exc_info: await OrgStore.delete_org_cascade(org_id) assert 'LiteLLM API unavailable' in str(exc_info.value) # Verify transaction was rolled back - organization should still exist async with async_session_maker() as session: persisted_org = await session.get(Org, org_id) assert persisted_org is not None assert persisted_org.name == 'Test Organization' # Org member should still exist result = await session.execute(select(OrgMember).filter_by(org_id=org_id)) persisted_member = result.scalars().first() assert persisted_member is not None @pytest.mark.asyncio async def test_get_user_orgs_paginated_first_page( async_session_maker, mock_litellm_api ): """ GIVEN: User is member of multiple organizations WHEN: get_user_orgs_paginated is called without page_id THEN: First page of organizations is returned in alphabetical order """ # Arrange user_id = uuid.uuid4() other_user_id = uuid.uuid4() async with async_session_maker() as session: # Create orgs for the user org1 = Org(name='Alpha Org') org2 = Org(name='Beta Org') org3 = Org(name='Gamma Org') # Create org for another user (should not be included) org4 = Org(name='Other Org') session.add_all([org1, org2, org3, org4]) await session.flush() # Create user and role user = User(id=user_id, current_org_id=org1.id) other_user = User(id=other_user_id, current_org_id=org4.id) role = Role(id=1, name='member', rank=2) session.add_all([user, other_user, role]) await session.flush() # Create memberships member1 = OrgMember( org_id=org1.id, user_id=user_id, role_id=1, llm_api_key='key1' ) member2 = OrgMember( org_id=org2.id, user_id=user_id, role_id=1, llm_api_key='key2' ) member3 = OrgMember( org_id=org3.id, user_id=user_id, role_id=1, llm_api_key='key3' ) other_member = OrgMember( org_id=org4.id, user_id=other_user_id, role_id=1, llm_api_key='key4' ) session.add_all([member1, member2, member3, other_member]) await session.commit() # Act with patch('storage.org_store.a_session_maker', async_session_maker): orgs, next_page_id = await OrgStore.get_user_orgs_paginated( user_id=user_id, page_id=None, limit=2 ) # Assert assert len(orgs) == 2 assert orgs[0].name == 'Alpha Org' assert orgs[1].name == 'Beta Org' assert next_page_id == '2' # Has more results # Verify other user's org is not included org_names = [org.name for org in orgs] assert 'Other Org' not in org_names @pytest.mark.asyncio async def test_get_user_orgs_paginated_with_page_id( async_session_maker, mock_litellm_api ): """ GIVEN: User has multiple organizations and page_id is provided WHEN: get_user_orgs_paginated is called with page_id THEN: Organizations starting from offset are returned """ # Arrange user_id = uuid.uuid4() async with async_session_maker() as session: org1 = Org(name='Alpha Org') org2 = Org(name='Beta Org') org3 = Org(name='Gamma Org') session.add_all([org1, org2, org3]) await session.flush() user = User(id=user_id, current_org_id=org1.id) role = Role(id=1, name='member', rank=2) session.add_all([user, role]) await session.flush() member1 = OrgMember( org_id=org1.id, user_id=user_id, role_id=1, llm_api_key='key1' ) member2 = OrgMember( org_id=org2.id, user_id=user_id, role_id=1, llm_api_key='key2' ) member3 = OrgMember( org_id=org3.id, user_id=user_id, role_id=1, llm_api_key='key3' ) session.add_all([member1, member2, member3]) await session.commit() # Act with patch('storage.org_store.a_session_maker', async_session_maker): orgs, next_page_id = await OrgStore.get_user_orgs_paginated( user_id=user_id, page_id='1', limit=1 ) # Assert assert len(orgs) == 1 assert orgs[0].name == 'Beta Org' # Second org (offset 1) assert next_page_id == '2' # Has more results @pytest.mark.asyncio async def test_get_user_orgs_paginated_no_more_results( async_session_maker, mock_litellm_api ): """ GIVEN: User has organizations but fewer than limit WHEN: get_user_orgs_paginated is called THEN: All organizations are returned and next_page_id is None """ # Arrange user_id = uuid.uuid4() async with async_session_maker() as session: org1 = Org(name='Alpha Org') org2 = Org(name='Beta Org') session.add_all([org1, org2]) await session.flush() user = User(id=user_id, current_org_id=org1.id) role = Role(id=1, name='member', rank=2) session.add_all([user, role]) await session.flush() member1 = OrgMember( org_id=org1.id, user_id=user_id, role_id=1, llm_api_key='key1' ) member2 = OrgMember( org_id=org2.id, user_id=user_id, role_id=1, llm_api_key='key2' ) session.add_all([member1, member2]) await session.commit() # Act with patch('storage.org_store.a_session_maker', async_session_maker): orgs, next_page_id = await OrgStore.get_user_orgs_paginated( user_id=user_id, page_id=None, limit=10 ) # Assert assert len(orgs) == 2 assert next_page_id is None @pytest.mark.asyncio async def test_get_user_orgs_paginated_invalid_page_id( async_session_maker, mock_litellm_api ): """ GIVEN: Invalid page_id (non-numeric string) WHEN: get_user_orgs_paginated is called THEN: Results start from beginning (offset 0) """ # Arrange user_id = uuid.uuid4() async with async_session_maker() as session: org1 = Org(name='Alpha Org') session.add(org1) await session.flush() user = User(id=user_id, current_org_id=org1.id) role = Role(id=1, name='member', rank=2) session.add_all([user, role]) await session.flush() member1 = OrgMember( org_id=org1.id, user_id=user_id, role_id=1, llm_api_key='key1' ) session.add(member1) await session.commit() # Act with patch('storage.org_store.a_session_maker', async_session_maker): orgs, next_page_id = await OrgStore.get_user_orgs_paginated( user_id=user_id, page_id='invalid', limit=10 ) # Assert assert len(orgs) == 1 assert orgs[0].name == 'Alpha Org' assert next_page_id is None @pytest.mark.asyncio async def test_get_user_orgs_paginated_empty_results(async_session_maker): """ GIVEN: User has no organizations WHEN: get_user_orgs_paginated is called THEN: Empty list and None next_page_id are returned """ # Arrange user_id = uuid.uuid4() # Act with patch('storage.org_store.a_session_maker', async_session_maker): orgs, next_page_id = await OrgStore.get_user_orgs_paginated( user_id=user_id, page_id=None, limit=10 ) # Assert assert len(orgs) == 0 assert next_page_id is None @pytest.mark.asyncio async def test_get_user_orgs_paginated_ordering(async_session_maker, mock_litellm_api): """ GIVEN: User has organizations with different names WHEN: get_user_orgs_paginated is called THEN: Organizations are returned in alphabetical order by name """ # Arrange user_id = uuid.uuid4() async with async_session_maker() as session: # Create orgs in non-alphabetical order org3 = Org(name='Zebra Org') org1 = Org(name='Apple Org') org2 = Org(name='Banana Org') session.add_all([org3, org1, org2]) await session.flush() user = User(id=user_id, current_org_id=org1.id) role = Role(id=1, name='member', rank=2) session.add_all([user, role]) await session.flush() member1 = OrgMember( org_id=org1.id, user_id=user_id, role_id=1, llm_api_key='key1' ) member2 = OrgMember( org_id=org2.id, user_id=user_id, role_id=1, llm_api_key='key2' ) member3 = OrgMember( org_id=org3.id, user_id=user_id, role_id=1, llm_api_key='key3' ) session.add_all([member1, member2, member3]) await session.commit() # Act with patch('storage.org_store.a_session_maker', async_session_maker): orgs, _ = await OrgStore.get_user_orgs_paginated( user_id=user_id, page_id=None, limit=10 ) # Assert assert len(orgs) == 3 assert orgs[0].name == 'Apple Org' assert orgs[1].name == 'Banana Org' assert orgs[2].name == 'Zebra Org' def test_orphaned_user_error_contains_user_ids(): """ GIVEN: OrphanedUserError is created with a list of user IDs WHEN: The error message is accessed THEN: Message includes the count and stores user IDs """ # Arrange from server.routes.org_models import OrphanedUserError user_ids = [str(uuid.uuid4()), str(uuid.uuid4())] # Act error = OrphanedUserError(user_ids) # Assert assert error.user_ids == user_ids assert '2 user(s)' in str(error) assert 'no remaining organization' in str(error) def test_org_deletion_with_invitations_uses_passive_deletes( session_maker, mock_litellm_api ): """ GIVEN: Organization has associated invitations with non-nullable org_id foreign key WHEN: Organization is deleted via SQLAlchemy session.delete() THEN: Deletion succeeds without NOT NULL constraint violation (passive_deletes=True defers to database CASCADE instead of setting org_id to NULL) This test verifies the fix for the bug where SQLAlchemy would try to SET org_id=NULL on org_invitation before deleting the org, causing: "NOT NULL constraint failed: org_invitation.org_id" With passive_deletes=True on the relationship, SQLAlchemy defers to the database's CASCADE constraint instead of trying to nullify the foreign key. Note: SQLite doesn't enforce CASCADE by default, so we only verify that the deletion succeeds. In production (PostgreSQL), CASCADE handles cleanup. """ from datetime import datetime, timedelta # Arrange org_id = uuid.uuid4() other_org_id = uuid.uuid4() user_id = uuid.uuid4() with session_maker() as session: # Create role first (required for invitation) role = Role(id=1, name='owner', rank=1) session.add(role) session.flush() # Create organization to be deleted org = Org(id=org_id, name='test-org-with-invitations') session.add(org) session.flush() # Create a second org for the user's current_org_id # (to avoid the user.current_org_id constraint issue during deletion) other_org = Org(id=other_org_id, name='other-org') session.add(other_org) session.flush() # Create user with current_org pointing to the OTHER org (not the one being deleted) user = User(id=user_id, current_org_id=other_org_id) session.add(user) session.flush() # Create invitation associated with the organization to be deleted invitation = OrgInvitation( token='test-invitation-token-12345', org_id=org_id, email='invitee@example.com', role_id=1, inviter_id=user_id, status='pending', created_at=datetime.now(), expires_at=datetime.now() + timedelta(days=7), ) session.add(invitation) session.commit() # Verify invitation was created invitation_count = session.query(OrgInvitation).filter_by(org_id=org_id).count() assert invitation_count == 1 # Act - Delete organization via SQLAlchemy (this is what triggered the bug) # Without passive_deletes=True, SQLAlchemy would try to SET org_id=NULL # which violates the NOT NULL constraint on org_invitation.org_id with session_maker() as session: org = session.query(Org).filter(Org.id == org_id).first() assert org is not None # This should NOT raise IntegrityError with passive_deletes=True # Previously this would fail with: # "NOT NULL constraint failed: org_invitation.org_id" session.delete(org) session.commit() # Success indicates passive_deletes=True is working # Assert - Organization should be deleted with session_maker() as session: deleted_org = session.query(Org).filter(Org.id == org_id).first() assert deleted_org is None # ============================================================================= # Tests for async organization-defaults propagation methods # ============================================================================= @pytest.mark.asyncio async def test_update_org_defaults_async_with_llm_api_key(): """GIVEN: Organization with members and llm_api_key in update settings WHEN: update_org_defaults_async is called THEN: Org fields are updated and llm_api_key is propagated to all members """ from server.routes.org_models import OrgUpdate # Arrange org_id = uuid.uuid4() mock_org = Org( id=org_id, name='Test Organization', agent_settings=OpenHandsAgentSettings(llm={'model': 'old-model'}), ) llm_settings = OrgUpdate( agent_settings_diff={'llm': {'model': 'new-model'}}, llm_api_key='new-member-api-key', ) # Mock the async session and member store mock_session = AsyncMock() mock_result = MagicMock() mock_result.scalars.return_value.first.return_value = mock_org mock_session.execute.return_value = mock_result mock_session.commit = AsyncMock() mock_session.refresh = AsyncMock() @asynccontextmanager async def mock_a_session_maker(): yield mock_session with ( patch('storage.org_store.a_session_maker', mock_a_session_maker), patch( 'storage.org_member_store.OrgMemberStore.update_all_members_settings_async', AsyncMock(), ) as mock_member_update, ): # Act result = await OrgStore.update_org_defaults_async( org_id, llm_settings, str(uuid.uuid4()), ) # Assert - Org is returned assert result is not None assert result.agent_settings['llm']['model'] == 'new-model' # Assert - Member update was called with correct settings mock_member_update.assert_called_once() call_args = mock_member_update.call_args member_settings = call_args[0][2] # Third positional arg is member_settings assert member_settings.llm_api_key.get_secret_value() == 'new-member-api-key' assert member_settings.agent_settings_diff == {'llm': {'model': 'new-model'}} @pytest.mark.asyncio async def test_update_org_defaults_async_propagates_managed_key_reset(): """GIVEN: A unified OrgUpdate save that resolves to a managed org key WHEN: update_org_defaults_async is called THEN: the propagated member update carries that key and resets the custom-key flag """ from server.routes.org_models import OrgUpdate org_id = uuid.uuid4() user_id = str(uuid.uuid4()) mock_org = Org( id=org_id, name='Test Organization', agent_settings=OpenHandsAgentSettings(llm={'model': 'openhands/claude-3'}), ) update_data = OrgUpdate( agent_settings_diff={'llm': {'model': 'openhands/claude-3'}} ) mock_session = AsyncMock() mock_result = MagicMock() mock_result.scalars.return_value.first.return_value = mock_org mock_session.execute.return_value = mock_result mock_session.commit = AsyncMock() mock_session.refresh = AsyncMock() @asynccontextmanager async def mock_a_session_maker(): yield mock_session with ( patch('storage.org_store.a_session_maker', mock_a_session_maker), patch( 'storage.org_store.OrgStore._maybe_get_managed_llm_key_for_user', AsyncMock(return_value='managed-key'), ), patch( 'storage.org_member_store.OrgMemberStore.update_all_members_settings_async', AsyncMock(), ) as mock_member_update, ): result = await OrgStore.update_org_defaults_async(org_id, update_data, user_id) assert result is not None agent_settings = OrgStore.get_agent_settings_from_org(result) assert agent_settings.llm.model == 'litellm_proxy/claude-3' mock_member_update.assert_called_once() member_settings = mock_member_update.call_args[0][2] assert member_settings.llm_api_key.get_secret_value() == 'managed-key' assert member_settings.has_custom_llm_api_key is False @pytest.mark.asyncio async def test_update_org_defaults_async_non_key_changes_keep_custom_key_flags(): """GIVEN: An org-defaults save that only updates shared settings WHEN: update_org_defaults_async is called THEN: member propagation keeps personal custom-key flags untouched """ from server.routes.org_models import OrgUpdate org_id = uuid.uuid4() user_id = str(uuid.uuid4()) mock_org = Org( id=org_id, name='Test Organization', agent_settings=OpenHandsAgentSettings(llm={'model': 'openhands/claude-3'}), conversation_settings=ConversationSettings(), ) update_data = OrgUpdate(conversation_settings_diff={'max_iterations': 42}) mock_session = AsyncMock() mock_result = MagicMock() mock_result.scalars.return_value.first.return_value = mock_org mock_session.execute.return_value = mock_result mock_session.commit = AsyncMock() mock_session.refresh = AsyncMock() @asynccontextmanager async def mock_a_session_maker(): yield mock_session with ( patch('storage.org_store.a_session_maker', mock_a_session_maker), patch( 'storage.org_store.OrgStore._maybe_get_managed_llm_key_for_user', AsyncMock(return_value=None), ), patch( 'storage.org_member_store.OrgMemberStore.update_all_members_settings_async', AsyncMock(), ) as mock_member_update, ): await OrgStore.update_org_defaults_async(org_id, update_data, user_id) mock_member_update.assert_called_once() member_settings = mock_member_update.call_args[0][2] assert member_settings.conversation_settings_diff == {'max_iterations': 42} assert member_settings.has_custom_llm_api_key is None @pytest.mark.asyncio async def test_update_org_defaults_async_org_not_found(): """GIVEN: Non-existent organization ID WHEN: update_org_defaults_async is called THEN: Returns None """ from server.routes.org_models import OrgUpdate # Arrange non_existent_org_id = uuid.uuid4() llm_settings = OrgUpdate(agent_settings_diff={'llm': {'model': 'new-model'}}) # Mock the async session to return None for org mock_session = AsyncMock() mock_result = MagicMock() mock_result.scalars.return_value.first.return_value = None mock_session.execute.return_value = mock_result @asynccontextmanager async def mock_a_session_maker(): yield mock_session # Act with patch('storage.org_store.a_session_maker', mock_a_session_maker): result = await OrgStore.update_org_defaults_async( non_existent_org_id, llm_settings, str(uuid.uuid4()), ) # Assert assert result is None