Files
OpenHands/enterprise/tests/unit/test_org_store.py
2026-03-13 22:39:07 +00:00

1106 lines
35 KiB
Python

import uuid
from contextlib import asynccontextmanager
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from pydantic import SecretStr
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.storage.data_models.settings import Settings
@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='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),
):
updated_org = await OrgStore.update_org(
org_id=org_id, kwargs={'name': 'updated-org', 'agent': 'PlannerAgent'}
)
assert updated_org is not None
assert updated_org.name == 'updated-org'
assert updated_org.agent == 'PlannerAgent'
@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(), kwargs={'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': 'CodeActAgent'}
)
assert org is not None
assert org.name == 'new-org'
assert org.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(
language='es',
agent='CodeActAgent',
llm_model='gpt-4',
llm_api_key=SecretStr('test-key'),
enable_sound_notifications=True,
)
kwargs = OrgStore.get_kwargs_from_settings(settings)
# Should only include fields that exist in Org model
assert 'agent' in kwargs
assert 'default_llm_model' in kwargs
assert kwargs['agent'] == 'CodeActAgent'
assert kwargs['default_llm_model'] == 'gpt-4'
# 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='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 == '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='CodeActAgent',
default_max_iterations=50,
confirmation_mode=True,
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',
max_iterations=100,
llm_model='gpt-4',
)
# 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 == 'CodeActAgent'
assert result.default_max_iterations == 50
assert result.confirmation_mode is True
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 == 'CodeActAgent'
assert persisted_org.default_max_iterations == 50
assert persisted_org.confirmation_mode is True
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.max_iterations == 100
assert persisted_member.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 LLM settings methods
# =============================================================================
@pytest.mark.asyncio
async def test_update_org_llm_settings_async_with_llm_api_key():
"""
GIVEN: Organization with members and llm_api_key in update settings
WHEN: update_org_llm_settings_async is called
THEN: Org fields are updated and llm_api_key is propagated to all members
"""
from server.routes.org_models import OrgLLMSettingsUpdate
# Arrange
org_id = uuid.uuid4()
mock_org = Org(
id=org_id,
name='Test Organization',
default_llm_model='old-model',
)
llm_settings = OrgLLMSettingsUpdate(
default_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_llm_settings_async',
AsyncMock(),
) as mock_member_update,
):
# Act
result = await OrgStore.update_org_llm_settings_async(org_id, llm_settings)
# Assert - Org is returned
assert result is not None
assert result.default_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 == 'new-member-api-key'
assert member_settings.llm_model == 'new-model'
@pytest.mark.asyncio
async def test_update_org_llm_settings_async_org_not_found():
"""
GIVEN: Non-existent organization ID
WHEN: update_org_llm_settings_async is called
THEN: Returns None
"""
from server.routes.org_models import OrgLLMSettingsUpdate
# Arrange
non_existent_org_id = uuid.uuid4()
llm_settings = OrgLLMSettingsUpdate(default_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_llm_settings_async(
non_existent_org_id, llm_settings
)
# Assert
assert result is None