From 1e6a92b454456bd970a705d7814f7be3862037d0 Mon Sep 17 00:00:00 2001 From: Hiep Le <69354317+hieptl@users.noreply.github.com> Date: Mon, 2 Mar 2026 00:06:37 +0700 Subject: [PATCH] feat(backend): organizations llm settings api (org project) (#13108) --- enterprise/server/routes/org_models.py | 109 ++++++ enterprise/server/routes/orgs.py | 108 ++++++ .../services/org_llm_settings_service.py | 130 +++++++ enterprise/storage/org_llm_settings_store.py | 83 +++++ enterprise/storage/org_member_store.py | 30 +- enterprise/storage/org_service.py | 94 ----- enterprise/storage/org_store.py | 50 ++- .../services/test_org_llm_settings_service.py | 215 ++++++++++++ .../storage/test_org_llm_settings_store.py | 180 ++++++++++ .../tests/unit/test_org_member_store.py | 326 ++++++++++++++++++ enterprise/tests/unit/test_org_service.py | 314 ----------------- enterprise/tests/unit/test_org_store.py | 96 ++++++ 12 files changed, 1323 insertions(+), 412 deletions(-) create mode 100644 enterprise/server/services/org_llm_settings_service.py create mode 100644 enterprise/storage/org_llm_settings_store.py create mode 100644 enterprise/tests/unit/server/services/test_org_llm_settings_service.py create mode 100644 enterprise/tests/unit/storage/test_org_llm_settings_store.py diff --git a/enterprise/server/routes/org_models.py b/enterprise/server/routes/org_models.py index 825d16db22..a3be71d856 100644 --- a/enterprise/server/routes/org_models.py +++ b/enterprise/server/routes/org_models.py @@ -259,6 +259,115 @@ class OrgUpdate(BaseModel): condenser_max_size: int | None = Field(default=None, ge=20) +class OrgLLMSettingsResponse(BaseModel): + """Response model for organization LLM settings.""" + + default_llm_model: str | None = None + default_llm_base_url: str | None = None + search_api_key: str | None = None # Masked in response + agent: str | None = None + confirmation_mode: bool | None = None + security_analyzer: str | None = None + enable_default_condenser: bool = True + condenser_max_size: int | None = None + default_max_iterations: int | None = None + + @staticmethod + def _mask_key(secret: SecretStr | None) -> str | None: + """Mask an API key, showing only last 4 characters.""" + if secret is None: + return None + raw = secret.get_secret_value() + if not raw: + return None + if len(raw) <= 4: + return '****' + return '****' + raw[-4:] + + @classmethod + def from_org(cls, org: Org) -> 'OrgLLMSettingsResponse': + """Create response from Org entity.""" + return cls( + default_llm_model=org.default_llm_model, + default_llm_base_url=org.default_llm_base_url, + search_api_key=cls._mask_key(org.search_api_key), + agent=org.agent, + confirmation_mode=org.confirmation_mode, + security_analyzer=org.security_analyzer, + enable_default_condenser=org.enable_default_condenser + if org.enable_default_condenser is not None + else True, + condenser_max_size=org.condenser_max_size, + default_max_iterations=org.default_max_iterations, + ) + + +class OrgMemberLLMSettings(BaseModel): + """LLM settings to propagate to organization members. + + Field names match OrgMember DB columns. + """ + + llm_model: str | None = None + llm_base_url: str | None = None + max_iterations: int | None = None + llm_api_key: str | None = None + + def has_updates(self) -> bool: + """Check if any field is set (not None).""" + return any(getattr(self, field) is not None for field in self.model_fields) + + +class OrgLLMSettingsUpdate(BaseModel): + """Request model for updating organization LLM settings. + + Field names match Org DB columns exactly. + """ + + default_llm_model: str | None = None + default_llm_base_url: str | None = None + search_api_key: str | None = None + agent: str | None = None + confirmation_mode: bool | None = None + security_analyzer: str | None = None + enable_default_condenser: bool | None = None + condenser_max_size: int | None = Field(default=None, ge=20) + default_max_iterations: int | None = Field(default=None, gt=0) + llm_api_key: str | None = None + + def has_updates(self) -> bool: + """Check if any field is set (not None).""" + return any(getattr(self, field) is not None for field in self.model_fields) + + def apply_to_org(self, org: Org) -> None: + """Apply non-None settings to the organization model. + + Args: + org: Organization entity to update in place + """ + for field_name in self.model_fields: + value = getattr(self, field_name) + # Skip llm_api_key - it's only for member propagation, not org-level + if value is not None and field_name != 'llm_api_key': + setattr(org, field_name, value) + + def get_member_updates(self) -> OrgMemberLLMSettings | None: + """Get updates that need to be propagated to org members. + + Returns: + OrgMemberLLMSettings with mapped field values, or None if no member updates needed. + Maps: default_llm_model → llm_model, default_llm_base_url → llm_base_url, + default_max_iterations → max_iterations, llm_api_key → llm_api_key + """ + member_settings = OrgMemberLLMSettings( + llm_model=self.default_llm_model, + llm_base_url=self.default_llm_base_url, + max_iterations=self.default_max_iterations, + llm_api_key=self.llm_api_key, + ) + return member_settings if member_settings.has_updates() else None + + class OrgMemberResponse(BaseModel): """Response model for a single organization member.""" diff --git a/enterprise/server/routes/orgs.py b/enterprise/server/routes/orgs.py index 338b2b2ea0..6309ba1bb7 100644 --- a/enterprise/server/routes/orgs.py +++ b/enterprise/server/routes/orgs.py @@ -20,6 +20,8 @@ from server.routes.org_models import ( OrgAuthorizationError, OrgCreate, OrgDatabaseError, + OrgLLMSettingsResponse, + OrgLLMSettingsUpdate, OrgMemberNotFoundError, OrgMemberPage, OrgMemberResponse, @@ -36,6 +38,10 @@ from server.services.org_app_settings_service import ( OrgAppSettingsService, OrgAppSettingsServiceInjector, ) +from server.services.org_llm_settings_service import ( + OrgLLMSettingsService, + OrgLLMSettingsServiceInjector, +) from server.services.org_member_service import OrgMemberService from storage.org_service import OrgService from storage.user_store import UserStore @@ -46,6 +52,9 @@ from openhands.server.user_auth import get_user_id # Initialize API router org_router = APIRouter(prefix='/api/organizations', tags=['Orgs']) +# Create injector instance and dependency for LLM settings +_org_llm_settings_injector = OrgLLMSettingsServiceInjector() +org_llm_settings_service_dependency = Depends(_org_llm_settings_injector.depends) # Create injector instance and dependency at module level _org_app_settings_injector = OrgAppSettingsServiceInjector() org_app_settings_service_dependency = Depends(_org_app_settings_injector.depends) @@ -211,6 +220,105 @@ async def create_org( ) +@org_router.get( + '/llm', + response_model=OrgLLMSettingsResponse, + dependencies=[Depends(require_permission(Permission.VIEW_LLM_SETTINGS))], +) +async def get_org_llm_settings( + service: OrgLLMSettingsService = org_llm_settings_service_dependency, +) -> OrgLLMSettingsResponse: + """Get LLM settings for the user's current organization. + + This endpoint retrieves the LLM configuration settings for the + authenticated user's current organization. All organization members + can view these settings. + + Args: + service: OrgLLMSettingsService (injected by dependency) + + Returns: + OrgLLMSettingsResponse: The organization's LLM settings + + Raises: + HTTPException: 401 if not authenticated + HTTPException: 403 if not a member of any organization + HTTPException: 404 if current organization not found + HTTPException: 500 if retrieval fails + """ + try: + return await service.get_org_llm_settings() + except OrgNotFoundError as e: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=str(e), + ) + except Exception as e: + logger.exception( + 'Error getting organization LLM settings', + extra={'error': str(e)}, + ) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail='Failed to retrieve LLM settings', + ) + + +@org_router.post( + '/llm', + response_model=OrgLLMSettingsResponse, + dependencies=[Depends(require_permission(Permission.EDIT_LLM_SETTINGS))], +) +async def update_org_llm_settings( + settings: OrgLLMSettingsUpdate, + service: OrgLLMSettingsService = org_llm_settings_service_dependency, +) -> OrgLLMSettingsResponse: + """Update LLM settings for the user's current organization. + + This endpoint updates the LLM configuration settings for the + authenticated user's current organization. Only admins and owners + can update these settings. + + Args: + settings: The LLM settings to update (only non-None fields are updated) + service: OrgLLMSettingsService (injected by dependency) + + Returns: + OrgLLMSettingsResponse: The updated organization's LLM settings + + Raises: + HTTPException: 401 if not authenticated + HTTPException: 403 if user lacks EDIT_LLM_SETTINGS permission + HTTPException: 404 if current organization not found + HTTPException: 500 if update fails + """ + try: + return await service.update_org_llm_settings(settings) + except OrgNotFoundError as e: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=str(e), + ) + except OrgDatabaseError as e: + logger.error( + 'Database error updating LLM settings', + extra={'error': str(e)}, + ) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail='Failed to update LLM settings', + ) + except Exception as e: + logger.exception( + 'Error updating organization LLM settings', + extra={'error': str(e)}, + ) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail='Failed to update LLM settings', + ) + + @org_router.get( '/app', response_model=OrgAppSettingsResponse, diff --git a/enterprise/server/services/org_llm_settings_service.py b/enterprise/server/services/org_llm_settings_service.py new file mode 100644 index 0000000000..b6ca42d87d --- /dev/null +++ b/enterprise/server/services/org_llm_settings_service.py @@ -0,0 +1,130 @@ +"""Service class for managing organization LLM settings. + +Separates business logic from route handlers. +Uses dependency injection for db_session and user_context. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import AsyncGenerator + +from fastapi import Request +from server.routes.org_models import ( + OrgLLMSettingsResponse, + OrgLLMSettingsUpdate, + OrgNotFoundError, +) +from storage.org_llm_settings_store import OrgLLMSettingsStore + +from openhands.app_server.services.injector import Injector, InjectorState +from openhands.app_server.user.user_context import UserContext +from openhands.core.logger import openhands_logger as logger + + +@dataclass +class OrgLLMSettingsService: + """Service for org LLM settings with injected dependencies.""" + + store: OrgLLMSettingsStore + user_context: UserContext + + async def get_org_llm_settings(self) -> OrgLLMSettingsResponse: + """Get LLM settings for user's current organization. + + User ID is obtained from the injected user_context. + + Returns: + OrgLLMSettingsResponse: The organization's LLM settings + + Raises: + ValueError: If user is not authenticated + OrgNotFoundError: If current organization not found + """ + user_id = await self.user_context.get_user_id() + if not user_id: + raise ValueError('User is not authenticated') + + logger.info( + 'Getting organization LLM settings', + extra={'user_id': user_id}, + ) + + org = await self.store.get_current_org_by_user_id(user_id) + + if not org: + raise OrgNotFoundError('No current organization') + + return OrgLLMSettingsResponse.from_org(org) + + async def update_org_llm_settings( + self, + update_data: OrgLLMSettingsUpdate, + ) -> OrgLLMSettingsResponse: + """Update LLM settings for user's current organization. + + Only updates fields that are explicitly provided in update_data. + User ID is obtained from the injected user_context. + Session auto-commits at request end via DbSessionInjector. + + Args: + update_data: The update data from the request + + Returns: + OrgLLMSettingsResponse: The updated organization's LLM settings + + Raises: + ValueError: If user is not authenticated + OrgNotFoundError: If current organization not found + """ + user_id = await self.user_context.get_user_id() + if not user_id: + raise ValueError('User is not authenticated') + + logger.info( + 'Updating organization LLM settings', + extra={'user_id': user_id}, + ) + + # Check if any fields are provided + if not update_data.has_updates(): + # No fields to update, just return current settings + return await self.get_org_llm_settings() + + # Get user's current org first + org = await self.store.get_current_org_by_user_id(user_id) + if not org: + raise OrgNotFoundError('No current organization') + + # Update the org LLM settings + updated_org = await self.store.update_org_llm_settings( + org_id=org.id, + update_data=update_data, + ) + + if not updated_org: + raise OrgNotFoundError(str(org.id)) + + logger.info( + 'Organization LLM settings updated successfully', + extra={'user_id': user_id, 'org_id': str(org.id)}, + ) + + return OrgLLMSettingsResponse.from_org(updated_org) + + +class OrgLLMSettingsServiceInjector(Injector[OrgLLMSettingsService]): + """Injector that composes store and user_context for OrgLLMSettingsService.""" + + async def inject( + self, state: InjectorState, request: Request | None = None + ) -> AsyncGenerator[OrgLLMSettingsService, None]: + # Local imports to avoid circular dependencies + from openhands.app_server.config import get_db_session, get_user_context + + async with ( + get_user_context(state, request) as user_context, + get_db_session(state, request) as db_session, + ): + store = OrgLLMSettingsStore(db_session=db_session) + yield OrgLLMSettingsService(store=store, user_context=user_context) diff --git a/enterprise/storage/org_llm_settings_store.py b/enterprise/storage/org_llm_settings_store.py new file mode 100644 index 0000000000..ad901b4ad6 --- /dev/null +++ b/enterprise/storage/org_llm_settings_store.py @@ -0,0 +1,83 @@ +"""Store class for managing organization LLM settings.""" + +from __future__ import annotations + +import uuid +from dataclasses import dataclass +from uuid import UUID + +from server.routes.org_models import OrgLLMSettingsUpdate +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession +from storage.org import Org +from storage.org_member_store import OrgMemberStore +from storage.user import User + + +@dataclass +class OrgLLMSettingsStore: + """Store for org LLM settings with injected db_session.""" + + db_session: AsyncSession + + async def get_current_org_by_user_id(self, user_id: str) -> Org | None: + """Get the user's current organization. + + Args: + user_id: The user's ID (Keycloak user ID) + + Returns: + Org: The user's current organization, or None if not found + """ + # First get the user to find their current_org_id + result = await self.db_session.execute( + select(User).filter(User.id == uuid.UUID(user_id)) + ) + user = result.scalars().first() + + if not user or not user.current_org_id: + return None + + # Then get the org + result = await self.db_session.execute( + select(Org).filter(Org.id == user.current_org_id) + ) + return result.scalars().first() + + async def update_org_llm_settings( + self, org_id: UUID, update_data: OrgLLMSettingsUpdate + ) -> Org | None: + """Update organization LLM settings. + + Also propagates relevant settings to all org members. + Uses flush() - commit happens at request end via DbSessionInjector. + + Args: + org_id: The organization's ID + update_data: Pydantic model with fields to update + + Returns: + Org: The updated organization, or None if org not found + """ + result = await self.db_session.execute( + select(Org).filter(Org.id == org_id).with_for_update() + ) + org = result.scalars().first() + + if not org: + return None + + # Apply updates to org (excludes llm_api_key which is member-only) + update_data.apply_to_org(org) + + # Propagate relevant settings to all org members + member_updates = update_data.get_member_updates() + if member_updates: + await OrgMemberStore.update_all_members_llm_settings_async( + self.db_session, org_id, member_updates + ) + + # flush instead of commit - DbSessionInjector auto-commits at request end + await self.db_session.flush() + await self.db_session.refresh(org) + return org diff --git a/enterprise/storage/org_member_store.py b/enterprise/storage/org_member_store.py index e79b8a3a34..c92d7ba867 100644 --- a/enterprise/storage/org_member_store.py +++ b/enterprise/storage/org_member_store.py @@ -5,9 +5,12 @@ Store class for managing organization-member relationships. from typing import Optional from uuid import UUID -from sqlalchemy import func, select +from server.routes.org_models import OrgMemberLLMSettings +from sqlalchemy import func, select, update +from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import joinedload from storage.database import a_session_maker, session_maker +from storage.encrypt_utils import encrypt_value from storage.org_member import OrgMember from storage.user import User from storage.user_settings import UserSettings @@ -254,3 +257,28 @@ class OrgMemberStore: members = members[:limit] return members, has_more + + @staticmethod + async def update_all_members_llm_settings_async( + session: AsyncSession, + org_id: UUID, + member_settings: OrgMemberLLMSettings, + ) -> None: + """Update LLM settings for all members of an organization. + + Args: + session: Database session (passed from caller for transaction) + org_id: Organization ID + member_settings: Typed LLM settings to apply to all members + """ + # Build update values from non-None fields + values = member_settings.model_dump(exclude_none=True) + + # Handle encrypted llm_api_key field - map to _llm_api_key column with encryption + if 'llm_api_key' in values: + raw_key = values.pop('llm_api_key') + values['_llm_api_key'] = encrypt_value(raw_key) + + if values: + stmt = update(OrgMember).where(OrgMember.org_id == org_id).values(**values) + await session.execute(stmt) diff --git a/enterprise/storage/org_service.py b/enterprise/storage/org_service.py index 1489369b10..00fdf4443e 100644 --- a/enterprise/storage/org_service.py +++ b/enterprise/storage/org_service.py @@ -9,7 +9,6 @@ from uuid import UUID as parse_uuid from server.constants import ORG_SETTINGS_VERSION, get_default_litellm_model from server.routes.org_models import ( LiteLLMIntegrationError, - OrgAppSettingsUpdate, OrgAuthorizationError, OrgDatabaseError, OrgNameExistsError, @@ -952,96 +951,3 @@ class OrgService: extra={'user_id': user_id, 'org_id': str(org_id), 'error': str(e)}, ) raise OrgDatabaseError(f'Failed to switch organization: {str(e)}') - - @staticmethod - def get_current_org_app_settings(user_id: str) -> Org: - """ - Get organization app settings for the user's current organization. - - Args: - user_id: User ID - - Returns: - Org: The user's current organization - - Raises: - OrgNotFoundError: If user has no current organization - """ - logger.info( - 'Retrieving organization app settings', - extra={'user_id': user_id}, - ) - - org = OrgStore.get_current_org_from_keycloak_user_id(user_id) - if not org: - logger.warning( - 'Current organization not found for user', - extra={'user_id': user_id}, - ) - raise OrgNotFoundError('current') - - logger.info( - 'Successfully retrieved organization app settings', - extra={'user_id': user_id, 'org_id': str(org.id)}, - ) - - return org - - @staticmethod - def update_current_org_app_settings( - user_id: str, - update_data: OrgAppSettingsUpdate, - ) -> Org: - """ - Update organization app settings for the user's current organization. - - Args: - user_id: User ID - update_data: App settings update data - - Returns: - Org: The updated organization - - Raises: - OrgNotFoundError: If user has no current organization - OrgDatabaseError: If update fails - """ - logger.info( - 'Updating organization app settings', - extra={'user_id': user_id}, - ) - - org = OrgStore.get_current_org_from_keycloak_user_id(user_id) - if not org: - logger.warning( - 'Current organization not found for user', - extra={'user_id': user_id}, - ) - raise OrgNotFoundError('current') - - update_dict = update_data.model_dump(exclude_unset=True) - if not update_dict: - logger.info( - 'No fields to update in app settings', - extra={'user_id': user_id, 'org_id': str(org.id)}, - ) - return org - - updated_org = OrgStore.update_org(org.id, update_dict) - if not updated_org: - logger.error( - 'Failed to update organization app settings', - extra={'user_id': user_id, 'org_id': str(org.id)}, - ) - raise OrgDatabaseError('Failed to update organization app settings') - - logger.info( - 'Successfully updated organization app settings', - extra={ - 'user_id': user_id, - 'org_id': str(org.id), - 'updated_fields': list(update_dict.keys()), - }, - ) - - return updated_org diff --git a/enterprise/storage/org_store.py b/enterprise/storage/org_store.py index 153c2726e5..85b86ab934 100644 --- a/enterprise/storage/org_store.py +++ b/enterprise/storage/org_store.py @@ -10,10 +10,10 @@ from server.constants import ( ORG_SETTINGS_VERSION, get_default_litellm_model, ) -from server.routes.org_models import OrphanedUserError -from sqlalchemy import text +from server.routes.org_models import OrgLLMSettingsUpdate, OrphanedUserError +from sqlalchemy import select, text from sqlalchemy.orm import joinedload -from storage.database import session_maker +from storage.database import a_session_maker, session_maker from storage.lite_llm_manager import LiteLlmManager from storage.org import Org from storage.org_member import OrgMember @@ -386,3 +386,47 @@ class OrgStore: extra={'org_id': str(org_id), 'error': str(e)}, ) raise + + @staticmethod + async def get_org_by_id_async(org_id: UUID) -> Org | None: + """Get organization by ID (async version).""" + async with a_session_maker() as session: + result = await session.execute(select(Org).filter(Org.id == org_id)) + org = result.scalars().first() + return OrgStore._validate_org_version(org) if org else None + + @staticmethod + async def update_org_llm_settings_async( + org_id: UUID, + llm_settings: OrgLLMSettingsUpdate, + ) -> Org | None: + """Update organization LLM settings and propagate to members (async version). + + Args: + org_id: Organization ID + llm_settings: Typed LLM settings update model + + Returns: + Updated Org or None if not found + """ + from storage.org_member_store import OrgMemberStore + + async with a_session_maker() as session: + result = await session.execute(select(Org).filter(Org.id == org_id)) + org = result.scalars().first() + if not org: + return None + + # Apply updates to org + llm_settings.apply_to_org(org) + + # Propagate relevant settings to all org members + member_updates = llm_settings.get_member_updates() + if member_updates: + await OrgMemberStore.update_all_members_llm_settings_async( + session, org_id, member_updates + ) + + await session.commit() + await session.refresh(org) + return org diff --git a/enterprise/tests/unit/server/services/test_org_llm_settings_service.py b/enterprise/tests/unit/server/services/test_org_llm_settings_service.py new file mode 100644 index 0000000000..70c44d36f8 --- /dev/null +++ b/enterprise/tests/unit/server/services/test_org_llm_settings_service.py @@ -0,0 +1,215 @@ +""" +Unit tests for OrgLLMSettingsService. + +Tests the service layer for organization LLM settings operations. +""" + +import uuid +from unittest.mock import AsyncMock, MagicMock + +import pytest +from server.routes.org_models import ( + OrgLLMSettingsResponse, + OrgLLMSettingsUpdate, + OrgNotFoundError, +) +from server.services.org_llm_settings_service import OrgLLMSettingsService +from storage.org import Org + + +@pytest.fixture +def user_id(): + """Create a test user ID.""" + return str(uuid.uuid4()) + + +@pytest.fixture +def org_id(): + """Create a test org ID.""" + return uuid.uuid4() + + +@pytest.fixture +def mock_org(org_id): + """Create a mock organization with LLM settings.""" + org = MagicMock(spec=Org) + org.id = org_id + org.default_llm_model = 'claude-3' + org.default_llm_base_url = 'https://api.anthropic.com' + org.search_api_key = None + org.agent = 'CodeActAgent' + org.confirmation_mode = True + org.security_analyzer = None + org.enable_default_condenser = True + org.condenser_max_size = None + org.default_max_iterations = 50 + return org + + +@pytest.fixture +def mock_store(): + """Create a mock OrgLLMSettingsStore.""" + return MagicMock() + + +@pytest.fixture +def mock_user_context(user_id): + """Create a mock UserContext that returns the user_id.""" + context = MagicMock() + context.get_user_id = AsyncMock(return_value=user_id) + return context + + +@pytest.mark.asyncio +async def test_get_org_llm_settings_success( + user_id, mock_org, mock_store, mock_user_context +): + """ + GIVEN: A user with a current organization + WHEN: get_org_llm_settings is called + THEN: OrgLLMSettingsResponse is returned with correct data + """ + # Arrange + mock_store.get_current_org_by_user_id = AsyncMock(return_value=mock_org) + service = OrgLLMSettingsService(store=mock_store, user_context=mock_user_context) + + # Act + result = await service.get_org_llm_settings() + + # Assert + assert isinstance(result, OrgLLMSettingsResponse) + assert result.default_llm_model == 'claude-3' + assert result.agent == 'CodeActAgent' + mock_store.get_current_org_by_user_id.assert_called_once_with(user_id) + + +@pytest.mark.asyncio +async def test_get_org_llm_settings_user_not_authenticated(mock_store): + """ + GIVEN: A user is not authenticated + WHEN: get_org_llm_settings is called + THEN: ValueError is raised + """ + # Arrange + mock_user_context = MagicMock() + mock_user_context.get_user_id = AsyncMock(return_value=None) + service = OrgLLMSettingsService(store=mock_store, user_context=mock_user_context) + + # Act & Assert + with pytest.raises(ValueError) as exc_info: + await service.get_org_llm_settings() + + assert 'not authenticated' in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_get_org_llm_settings_org_not_found( + user_id, mock_store, mock_user_context +): + """ + GIVEN: A user has no current organization + WHEN: get_org_llm_settings is called + THEN: OrgNotFoundError is raised + """ + # Arrange + mock_store.get_current_org_by_user_id = AsyncMock(return_value=None) + service = OrgLLMSettingsService(store=mock_store, user_context=mock_user_context) + + # Act & Assert + with pytest.raises(OrgNotFoundError) as exc_info: + await service.get_org_llm_settings() + + assert 'No current organization' in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_update_org_llm_settings_success( + user_id, mock_org, mock_store, mock_user_context +): + """ + GIVEN: A user with a current organization + WHEN: update_org_llm_settings is called with new values + THEN: OrgLLMSettingsResponse is returned with updated data + """ + # Arrange + updated_org = MagicMock(spec=Org) + updated_org.id = mock_org.id + updated_org.default_llm_model = 'new-model' + updated_org.default_llm_base_url = None + updated_org.search_api_key = None + updated_org.agent = 'CodeActAgent' + updated_org.confirmation_mode = False + updated_org.security_analyzer = None + updated_org.enable_default_condenser = True + updated_org.condenser_max_size = None + updated_org.default_max_iterations = 100 + + update_data = OrgLLMSettingsUpdate( + default_llm_model='new-model', + confirmation_mode=False, + default_max_iterations=100, + ) + + mock_store.get_current_org_by_user_id = AsyncMock(return_value=mock_org) + mock_store.update_org_llm_settings = AsyncMock(return_value=updated_org) + service = OrgLLMSettingsService(store=mock_store, user_context=mock_user_context) + + # Act + result = await service.update_org_llm_settings(update_data) + + # Assert + assert isinstance(result, OrgLLMSettingsResponse) + assert result.default_llm_model == 'new-model' + assert result.confirmation_mode is False + assert result.default_max_iterations == 100 + mock_store.update_org_llm_settings.assert_called_once_with( + org_id=mock_org.id, + update_data=update_data, + ) + + +@pytest.mark.asyncio +async def test_update_org_llm_settings_no_changes( + user_id, mock_org, mock_store, mock_user_context +): + """ + GIVEN: A user with a current organization + WHEN: update_org_llm_settings is called with no fields + THEN: Current settings are returned without calling update + """ + # Arrange + update_data = OrgLLMSettingsUpdate() # No fields set + + mock_store.get_current_org_by_user_id = AsyncMock(return_value=mock_org) + mock_store.update_org_llm_settings = AsyncMock() + service = OrgLLMSettingsService(store=mock_store, user_context=mock_user_context) + + # Act + result = await service.update_org_llm_settings(update_data) + + # Assert + assert isinstance(result, OrgLLMSettingsResponse) + assert result.default_llm_model == 'claude-3' + mock_store.update_org_llm_settings.assert_not_called() + + +@pytest.mark.asyncio +async def test_update_org_llm_settings_org_not_found( + user_id, mock_store, mock_user_context +): + """ + GIVEN: A user has no current organization + WHEN: update_org_llm_settings is called + THEN: OrgNotFoundError is raised + """ + # Arrange + update_data = OrgLLMSettingsUpdate(default_llm_model='new-model') + + mock_store.get_current_org_by_user_id = AsyncMock(return_value=None) + service = OrgLLMSettingsService(store=mock_store, user_context=mock_user_context) + + # Act & Assert + with pytest.raises(OrgNotFoundError) as exc_info: + await service.update_org_llm_settings(update_data) + + assert 'No current organization' in str(exc_info.value) diff --git a/enterprise/tests/unit/storage/test_org_llm_settings_store.py b/enterprise/tests/unit/storage/test_org_llm_settings_store.py new file mode 100644 index 0000000000..65fa19c816 --- /dev/null +++ b/enterprise/tests/unit/storage/test_org_llm_settings_store.py @@ -0,0 +1,180 @@ +""" +Unit tests for OrgLLMSettingsStore. + +Tests the async database operations for organization LLM settings. +""" + +import uuid +from unittest.mock import AsyncMock, patch + +import pytest +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine +from sqlalchemy.pool import StaticPool + +# Mock the database module before importing +with patch('storage.database.engine', create=True), patch( + 'storage.database.a_engine', create=True +): + from server.routes.org_models import OrgLLMSettingsUpdate + from storage.base import Base + from storage.org import Org + from storage.org_llm_settings_store import OrgLLMSettingsStore + from storage.user import User + + +@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, + ) + 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_maker(async_engine): + """Create an async session maker for testing.""" + return async_sessionmaker(async_engine, class_=AsyncSession, expire_on_commit=False) + + +@pytest.mark.asyncio +async def test_get_current_org_by_user_id_success(async_session_maker): + """ + GIVEN: A user exists with a current_org_id + WHEN: get_current_org_by_user_id is called + THEN: The user's current organization is returned + """ + # Arrange + async with async_session_maker() as session: + org = Org(name='test-org', default_llm_model='claude-3') + session.add(org) + await session.flush() + + user = User(id=uuid.uuid4(), current_org_id=org.id) + session.add(user) + await session.commit() + user_id = str(user.id) + + # Act + store = OrgLLMSettingsStore(db_session=session) + result = await store.get_current_org_by_user_id(user_id) + + # Assert + assert result is not None + assert result.name == 'test-org' + assert result.default_llm_model == 'claude-3' + + +@pytest.mark.asyncio +async def test_get_current_org_by_user_id_user_not_found(async_session_maker): + """ + GIVEN: A user does not exist in the database + WHEN: get_current_org_by_user_id is called + THEN: None is returned + """ + # Arrange + non_existent_id = str(uuid.uuid4()) + + # Act + async with async_session_maker() as session: + store = OrgLLMSettingsStore(db_session=session) + result = await store.get_current_org_by_user_id(non_existent_id) + + # Assert + assert result is None + + +@pytest.mark.asyncio +async def test_update_org_llm_settings_success(async_session_maker): + """ + GIVEN: An organization exists in the database + WHEN: update_org_llm_settings is called with new values + THEN: The organization's LLM settings are updated and returned + """ + # Arrange + async with async_session_maker() as session: + org = Org(name='test-org', default_llm_model='old-model') + session.add(org) + await session.commit() + org_id = org.id + + update_data = OrgLLMSettingsUpdate( + default_llm_model='new-model', + agent='CodeActAgent', + confirmation_mode=True, + ) + + # Act + store = OrgLLMSettingsStore(db_session=session) + with patch( + 'storage.org_llm_settings_store.OrgMemberStore.update_all_members_llm_settings_async', + AsyncMock(), + ): + result = await store.update_org_llm_settings(org_id, update_data) + + # Assert + assert result is not None + assert result.default_llm_model == 'new-model' + assert result.agent == 'CodeActAgent' + assert result.confirmation_mode is True + + +@pytest.mark.asyncio +async def test_update_org_llm_settings_org_not_found(async_session_maker): + """ + GIVEN: An organization does not exist in the database + WHEN: update_org_llm_settings is called + THEN: None is returned + """ + # Arrange + non_existent_org_id = uuid.uuid4() + update_data = OrgLLMSettingsUpdate(default_llm_model='new-model') + + # Act + async with async_session_maker() as session: + store = OrgLLMSettingsStore(db_session=session) + result = await store.update_org_llm_settings(non_existent_org_id, update_data) + + # Assert + assert result is None + + +@pytest.mark.asyncio +async def test_update_org_llm_settings_propagates_to_members(async_session_maker): + """ + GIVEN: An organization exists with update data containing member-relevant settings + WHEN: update_org_llm_settings is called + THEN: Member settings are propagated via OrgMemberStore + """ + # Arrange + async with async_session_maker() as session: + org = Org(name='test-org', default_llm_model='old-model') + session.add(org) + await session.commit() + org_id = org.id + + update_data = OrgLLMSettingsUpdate( + default_llm_model='new-model', + llm_api_key='new-api-key', + ) + + # Act + store = OrgLLMSettingsStore(db_session=session) + with patch( + 'storage.org_llm_settings_store.OrgMemberStore.update_all_members_llm_settings_async', + AsyncMock(), + ) as mock_update_members: + await store.update_org_llm_settings(org_id, update_data) + + # Assert + mock_update_members.assert_called_once() + call_args = mock_update_members.call_args + member_settings = call_args[0][2] + assert member_settings.llm_model == 'new-model' + assert member_settings.llm_api_key == 'new-api-key' diff --git a/enterprise/tests/unit/test_org_member_store.py b/enterprise/tests/unit/test_org_member_store.py index 5463d2d703..3f937c7f33 100644 --- a/enterprise/tests/unit/test_org_member_store.py +++ b/enterprise/tests/unit/test_org_member_store.py @@ -832,3 +832,329 @@ async def test_get_org_members_paginated_email_filter_case_insensitive( # Assert assert len(members) == 1 assert members[0].user.email == 'Alice@Example.COM' + + +@pytest.mark.asyncio +async def test_update_all_members_llm_settings_async_with_llm_api_key( + async_session_maker, +): + """ + GIVEN: Organization with members and llm_api_key in member settings + WHEN: update_all_members_llm_settings_async is called with llm_api_key + THEN: The llm_api_key is encrypted and stored in _llm_api_key column for all members + """ + from server.routes.org_models import OrgMemberLLMSettings + from storage.encrypt_utils import decrypt_value + + # Arrange + async with async_session_maker() as session: + org = Org(name='test-org') + session.add(org) + await session.flush() + + role = Role(name='member', rank=2) + session.add(role) + await session.flush() + + users = [ + User(id=uuid.uuid4(), current_org_id=org.id, email=f'user{i}@example.com') + for i in range(2) + ] + session.add_all(users) + await session.flush() + + org_members = [ + OrgMember( + org_id=org.id, + user_id=user.id, + role_id=role.id, + llm_api_key='old-key', + status='active', + ) + for user in users + ] + session.add_all(org_members) + await session.commit() + org_id = org.id + + # Act + new_api_key = 'new-test-api-key-12345' + member_settings = OrgMemberLLMSettings(llm_api_key=new_api_key) + + async with async_session_maker() as session: + await OrgMemberStore.update_all_members_llm_settings_async( + session, org_id, member_settings + ) + await session.commit() + + # Assert + async with async_session_maker() as session: + from sqlalchemy import select + + result = await session.execute( + select(OrgMember).filter(OrgMember.org_id == org_id) + ) + updated_members = result.scalars().all() + + assert len(updated_members) == 2 + for member in updated_members: + # Verify the encrypted value can be decrypted to the original + decrypted_key = decrypt_value(member._llm_api_key) + assert decrypted_key == new_api_key + + +@pytest.mark.asyncio +async def test_update_all_members_llm_settings_async_with_non_encrypted_fields( + async_session_maker, +): + """ + GIVEN: Organization with members + WHEN: update_all_members_llm_settings_async is called with non-encrypted fields + THEN: The fields are updated directly without encryption + """ + from server.routes.org_models import OrgMemberLLMSettings + + # Arrange + async with async_session_maker() as session: + org = Org(name='test-org') + session.add(org) + await session.flush() + + role = Role(name='member', rank=2) + session.add(role) + await session.flush() + + user = User(id=uuid.uuid4(), current_org_id=org.id, email='user@example.com') + session.add(user) + await session.flush() + + org_member = OrgMember( + org_id=org.id, + user_id=user.id, + role_id=role.id, + llm_api_key='test-key', + llm_model='old-model', + max_iterations=10, + status='active', + ) + session.add(org_member) + await session.commit() + org_id = org.id + + # Act + member_settings = OrgMemberLLMSettings( + llm_model='new-model', + llm_base_url='https://new-url.com', + max_iterations=50, + ) + + async with async_session_maker() as session: + await OrgMemberStore.update_all_members_llm_settings_async( + session, org_id, member_settings + ) + await session.commit() + + # Assert + async with async_session_maker() as session: + from sqlalchemy import select + + result = await session.execute( + select(OrgMember).filter(OrgMember.org_id == org_id) + ) + updated_member = result.scalars().first() + + assert updated_member.llm_model == 'new-model' + assert updated_member.llm_base_url == 'https://new-url.com' + assert updated_member.max_iterations == 50 + + +@pytest.mark.asyncio +async def test_update_all_members_llm_settings_async_with_empty_settings( + async_session_maker, +): + """ + GIVEN: Organization with members and empty member settings + WHEN: update_all_members_llm_settings_async is called with no fields set + THEN: No database update is performed + """ + from server.routes.org_models import OrgMemberLLMSettings + + # Arrange + async with async_session_maker() as session: + org = Org(name='test-org') + session.add(org) + await session.flush() + + role = Role(name='member', rank=2) + session.add(role) + await session.flush() + + user = User(id=uuid.uuid4(), current_org_id=org.id, email='user@example.com') + session.add(user) + await session.flush() + + org_member = OrgMember( + org_id=org.id, + user_id=user.id, + role_id=role.id, + llm_api_key='original-key', + llm_model='original-model', + status='active', + ) + session.add(org_member) + await session.commit() + org_id = org.id + + # Act - Empty settings (all None) + member_settings = OrgMemberLLMSettings() + + async with async_session_maker() as session: + await OrgMemberStore.update_all_members_llm_settings_async( + session, org_id, member_settings + ) + await session.commit() + + # Assert - Original values should be unchanged + async with async_session_maker() as session: + from sqlalchemy import select + + result = await session.execute( + select(OrgMember).filter(OrgMember.org_id == org_id) + ) + member = result.scalars().first() + + assert member.llm_model == 'original-model' + # Original key should still be there (encrypted) + assert member._llm_api_key is not None + + +# ============================================================================= +# OrgMemberLLMSettings and OrgLLMSettingsUpdate Model Unit Tests +# ============================================================================= + + +def test_org_member_llm_settings_has_updates_with_llm_api_key(): + """ + GIVEN: OrgMemberLLMSettings with only llm_api_key set + WHEN: has_updates() is called + THEN: Returns True + """ + from server.routes.org_models import OrgMemberLLMSettings + + # Arrange + settings = OrgMemberLLMSettings(llm_api_key='test-key') + + # Act + result = settings.has_updates() + + # Assert + assert result is True + + +def test_org_member_llm_settings_has_updates_empty(): + """ + GIVEN: OrgMemberLLMSettings with no fields set + WHEN: has_updates() is called + THEN: Returns False + """ + from server.routes.org_models import OrgMemberLLMSettings + + # Arrange + settings = OrgMemberLLMSettings() + + # Act + result = settings.has_updates() + + # Assert + assert result is False + + +def test_org_llm_settings_update_apply_to_org_skips_llm_api_key(): + """ + GIVEN: OrgLLMSettingsUpdate with llm_api_key and other fields set + WHEN: apply_to_org() is called + THEN: llm_api_key is NOT applied to org, but other fields are + """ + from unittest.mock import MagicMock + + from server.routes.org_models import OrgLLMSettingsUpdate + + # Arrange + settings = OrgLLMSettingsUpdate( + default_llm_model='claude-3', + llm_api_key='should-not-be-applied', + ) + mock_org = MagicMock() + mock_org.default_llm_model = None + + # Act + settings.apply_to_org(mock_org) + + # Assert + assert mock_org.default_llm_model == 'claude-3' + # llm_api_key should NOT be set on org (it's member-only) + assert ( + not hasattr(mock_org, 'llm_api_key') + or mock_org.llm_api_key != 'should-not-be-applied' + ) + + +def test_org_llm_settings_update_get_member_updates_includes_llm_api_key(): + """ + GIVEN: OrgLLMSettingsUpdate with llm_api_key set + WHEN: get_member_updates() is called + THEN: Returns OrgMemberLLMSettings with llm_api_key included + """ + from server.routes.org_models import OrgLLMSettingsUpdate + + # Arrange + settings = OrgLLMSettingsUpdate( + default_llm_model='claude-3', + llm_api_key='new-member-key', + ) + + # Act + member_updates = settings.get_member_updates() + + # Assert + assert member_updates is not None + assert member_updates.llm_api_key == 'new-member-key' + assert member_updates.llm_model == 'claude-3' + + +def test_org_llm_settings_update_get_member_updates_only_llm_api_key(): + """ + GIVEN: OrgLLMSettingsUpdate with only llm_api_key set + WHEN: get_member_updates() is called + THEN: Returns OrgMemberLLMSettings with llm_api_key (not None) + """ + from server.routes.org_models import OrgLLMSettingsUpdate + + # Arrange + settings = OrgLLMSettingsUpdate(llm_api_key='member-key-only') + + # Act + member_updates = settings.get_member_updates() + + # Assert + assert member_updates is not None + assert member_updates.llm_api_key == 'member-key-only' + assert member_updates.llm_model is None + + +def test_org_llm_settings_update_has_updates_with_llm_api_key(): + """ + GIVEN: OrgLLMSettingsUpdate with only llm_api_key set + WHEN: has_updates() is called + THEN: Returns True + """ + from server.routes.org_models import OrgLLMSettingsUpdate + + # Arrange + settings = OrgLLMSettingsUpdate(llm_api_key='test-key') + + # Act + result = settings.has_updates() + + # Assert + assert result is True diff --git a/enterprise/tests/unit/test_org_service.py b/enterprise/tests/unit/test_org_service.py index 61698ffb1e..47f7cd109a 100644 --- a/enterprise/tests/unit/test_org_service.py +++ b/enterprise/tests/unit/test_org_service.py @@ -2024,317 +2024,3 @@ async def test_switch_org_user_not_found(): await OrgService.switch_org(user_id, org_id) assert 'User not found' in str(exc_info.value) - - -# ============================================================================= -# Tests for App Settings Methods -# ============================================================================= - - -def test_get_current_org_app_settings_success(): - """ - GIVEN: User has a current organization - WHEN: get_current_org_app_settings is called - THEN: Organization is returned with app settings - """ - # Arrange - user_id = str(uuid.uuid4()) - org_id = uuid.uuid4() - mock_org = Org( - id=org_id, - name='Test Organization', - enable_proactive_conversation_starters=True, - enable_solvability_analysis=False, - max_budget_per_task=10.0, - ) - - with patch( - 'storage.org_service.OrgStore.get_current_org_from_keycloak_user_id', - return_value=mock_org, - ): - # Act - result = OrgService.get_current_org_app_settings(user_id) - - # Assert - assert result is not None - assert result.id == org_id - assert result.enable_proactive_conversation_starters is True - assert result.enable_solvability_analysis is False - assert result.max_budget_per_task == 10.0 - - -def test_get_current_org_app_settings_no_current_org(): - """ - GIVEN: User has no current organization - WHEN: get_current_org_app_settings is called - THEN: OrgNotFoundError is raised - """ - # Arrange - user_id = str(uuid.uuid4()) - - with patch( - 'storage.org_service.OrgStore.get_current_org_from_keycloak_user_id', - return_value=None, - ): - # Act & Assert - with pytest.raises(OrgNotFoundError) as exc_info: - OrgService.get_current_org_app_settings(user_id) - - assert 'current' in str(exc_info.value) - - -def test_update_current_org_app_settings_success(): - """ - GIVEN: User has a current organization and valid update data - WHEN: update_current_org_app_settings is called - THEN: Organization is updated and returned - """ - # Import here to avoid circular dependency in test setup - from server.routes.org_models import OrgAppSettingsUpdate - - # Arrange - user_id = str(uuid.uuid4()) - org_id = uuid.uuid4() - mock_org = Org( - id=org_id, - name='Test Organization', - enable_proactive_conversation_starters=True, - enable_solvability_analysis=None, - max_budget_per_task=None, - ) - updated_org = Org( - id=org_id, - name='Test Organization', - enable_proactive_conversation_starters=False, - enable_solvability_analysis=True, - max_budget_per_task=5.0, - ) - - update_data = OrgAppSettingsUpdate( - enable_proactive_conversation_starters=False, - enable_solvability_analysis=True, - max_budget_per_task=5.0, - ) - - with ( - patch( - 'storage.org_service.OrgStore.get_current_org_from_keycloak_user_id', - return_value=mock_org, - ), - patch( - 'storage.org_service.OrgStore.update_org', - return_value=updated_org, - ) as mock_update, - ): - # Act - result = OrgService.update_current_org_app_settings(user_id, update_data) - - # Assert - assert result is not None - assert result.enable_proactive_conversation_starters is False - assert result.enable_solvability_analysis is True - assert result.max_budget_per_task == 5.0 - mock_update.assert_called_once_with( - org_id, - { - 'enable_proactive_conversation_starters': False, - 'enable_solvability_analysis': True, - 'max_budget_per_task': 5.0, - }, - ) - - -def test_update_current_org_app_settings_partial_update(): - """ - GIVEN: User updates only some app settings - WHEN: update_current_org_app_settings is called - THEN: Only specified fields are updated - """ - from server.routes.org_models import OrgAppSettingsUpdate - - # Arrange - user_id = str(uuid.uuid4()) - org_id = uuid.uuid4() - mock_org = Org( - id=org_id, - name='Test Organization', - enable_proactive_conversation_starters=True, - enable_solvability_analysis=True, - max_budget_per_task=10.0, - ) - updated_org = Org( - id=org_id, - name='Test Organization', - enable_proactive_conversation_starters=False, - enable_solvability_analysis=True, - max_budget_per_task=10.0, - ) - - # Only update one field - update_data = OrgAppSettingsUpdate(enable_proactive_conversation_starters=False) - - with ( - patch( - 'storage.org_service.OrgStore.get_current_org_from_keycloak_user_id', - return_value=mock_org, - ), - patch( - 'storage.org_service.OrgStore.update_org', - return_value=updated_org, - ) as mock_update, - ): - # Act - result = OrgService.update_current_org_app_settings(user_id, update_data) - - # Assert - assert result is not None - assert result.enable_proactive_conversation_starters is False - # Only the specified field should be in the update dict - mock_update.assert_called_once_with( - org_id, - {'enable_proactive_conversation_starters': False}, - ) - - -def test_update_current_org_app_settings_empty_update(): - """ - GIVEN: No fields are specified in update - WHEN: update_current_org_app_settings is called - THEN: Original org is returned without database update - """ - from server.routes.org_models import OrgAppSettingsUpdate - - # Arrange - user_id = str(uuid.uuid4()) - org_id = uuid.uuid4() - mock_org = Org( - id=org_id, - name='Test Organization', - enable_proactive_conversation_starters=True, - enable_solvability_analysis=None, - max_budget_per_task=5.0, - ) - - update_data = OrgAppSettingsUpdate() # No fields set - - with ( - patch( - 'storage.org_service.OrgStore.get_current_org_from_keycloak_user_id', - return_value=mock_org, - ), - patch('storage.org_service.OrgStore.update_org') as mock_update, - ): - # Act - result = OrgService.update_current_org_app_settings(user_id, update_data) - - # Assert - assert result is mock_org - mock_update.assert_not_called() - - -def test_update_current_org_app_settings_set_to_null(): - """ - GIVEN: User explicitly sets max_budget_per_task to null - WHEN: update_current_org_app_settings is called - THEN: The field is set to null in the update - """ - from server.routes.org_models import OrgAppSettingsUpdate - - # Arrange - user_id = str(uuid.uuid4()) - org_id = uuid.uuid4() - mock_org = Org( - id=org_id, - name='Test Organization', - enable_proactive_conversation_starters=True, - enable_solvability_analysis=True, - max_budget_per_task=5.0, - ) - updated_org = Org( - id=org_id, - name='Test Organization', - enable_proactive_conversation_starters=True, - enable_solvability_analysis=True, - max_budget_per_task=None, - ) - - # Explicitly set max_budget_per_task to None - update_data = OrgAppSettingsUpdate.model_validate({'max_budget_per_task': None}) - - with ( - patch( - 'storage.org_service.OrgStore.get_current_org_from_keycloak_user_id', - return_value=mock_org, - ), - patch( - 'storage.org_service.OrgStore.update_org', - return_value=updated_org, - ) as mock_update, - ): - # Act - result = OrgService.update_current_org_app_settings(user_id, update_data) - - # Assert - assert result.max_budget_per_task is None - mock_update.assert_called_once_with( - org_id, - {'max_budget_per_task': None}, - ) - - -def test_update_current_org_app_settings_no_current_org(): - """ - GIVEN: User has no current organization - WHEN: update_current_org_app_settings is called - THEN: OrgNotFoundError is raised - """ - from server.routes.org_models import OrgAppSettingsUpdate - - # Arrange - user_id = str(uuid.uuid4()) - update_data = OrgAppSettingsUpdate(enable_proactive_conversation_starters=False) - - with patch( - 'storage.org_service.OrgStore.get_current_org_from_keycloak_user_id', - return_value=None, - ): - # Act & Assert - with pytest.raises(OrgNotFoundError) as exc_info: - OrgService.update_current_org_app_settings(user_id, update_data) - - assert 'current' in str(exc_info.value) - - -def test_update_current_org_app_settings_database_failure(): - """ - GIVEN: Database update fails - WHEN: update_current_org_app_settings is called - THEN: OrgDatabaseError is raised - """ - from server.routes.org_models import OrgAppSettingsUpdate - - # Arrange - user_id = str(uuid.uuid4()) - org_id = uuid.uuid4() - mock_org = Org( - id=org_id, - name='Test Organization', - enable_proactive_conversation_starters=True, - ) - update_data = OrgAppSettingsUpdate(enable_proactive_conversation_starters=False) - - with ( - patch( - 'storage.org_service.OrgStore.get_current_org_from_keycloak_user_id', - return_value=mock_org, - ), - patch( - 'storage.org_service.OrgStore.update_org', - return_value=None, - ), - ): - # Act & Assert - with pytest.raises(OrgDatabaseError) as exc_info: - OrgService.update_current_org_app_settings(user_id, update_data) - - assert 'Failed to update organization app settings' in str(exc_info.value) diff --git a/enterprise/tests/unit/test_org_store.py b/enterprise/tests/unit/test_org_store.py index 91ba4b60b7..6a31e35ff3 100644 --- a/enterprise/tests/unit/test_org_store.py +++ b/enterprise/tests/unit/test_org_store.py @@ -1,4 +1,5 @@ import uuid +from contextlib import asynccontextmanager from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -892,3 +893,98 @@ def test_org_deletion_with_invitations_uses_passive_deletes( 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