Spaces:
Sleeping
Sleeping
Download tests/database/test_core_models.py from techprotrade/annator-command-center: direct link, hf CLI and curl.
- Browser
- Download file 23.1 kB
-
https://huggingface.co/spaces/techprotrade/annator-command-center/resolve/main/tests/database/test_core_models.py
- Command line
-
hf download hf://spaces/techprotrade/annator-command-center/tests/database/test_core_models.py
-
curl -L -o test_core_models.py https://huggingface.co/spaces/techprotrade/annator-command-center/resolve/main/tests/database/test_core_models.py
23.1 kB
| """ | |
| Comprehensive core model tests covering CRUD, relationships, constraints, and cascade behaviors. | |
| Goal: Achieve 80%+ coverage for core models (Workspace, Team, Tenant, UserAccount, | |
| OAuthToken, ChatSession, ChatMessage) through comprehensive testing of: | |
| - CRUD operations (Create, Read, Update, Delete) | |
| - Relationship types (one-to-many, many-to-many, foreign keys) | |
| - Constraints (unique, not null, foreign keys) | |
| - Cascade behaviors (delete, nullify) | |
| - JSON field serialization | |
| - Timestamp auto-generation | |
| Tests use: | |
| - pytest fixtures for database sessions (db_session from conftest.py) | |
| - Factory pattern for test data creation (factories in tests/factories/) | |
| """ | |
| import pytest | |
| from datetime import datetime, timedelta | |
| from sqlalchemy.orm import Session | |
| from sqlalchemy.exc import IntegrityError | |
| from tests.factories.user_factory import UserFactory, AdminUserFactory | |
| from tests.factories.workspace_factory import WorkspaceFactory, TeamFactory | |
| from tests.factories.core_factory import ( | |
| TenantFactory, | |
| UserAccountFactory, | |
| OAuthTokenFactory, | |
| ChatMessageFactory, | |
| ) | |
| from tests.factories.chat_session_factory import ChatSessionFactory | |
| from core.models import ( | |
| Workspace, | |
| Team, | |
| Tenant, | |
| UserAccount, | |
| OAuthToken, | |
| ChatSession, | |
| ChatMessage, | |
| User, | |
| WorkspaceStatus, | |
| PlanType, | |
| ) | |
| # ============================================================================ | |
| # Task 2: Workspace Models Tests | |
| # ============================================================================ | |
| class TestWorkspaceModels: | |
| """Test Workspace model CRUD, relationships, and properties.""" | |
| def test_workspace_create_with_defaults(self, db_session: Session): | |
| """Test Workspace creation with minimal required fields.""" | |
| workspace = Workspace( | |
| name="Test Workspace", | |
| ) | |
| db_session.add(workspace) | |
| db_session.commit() | |
| db_session.refresh(workspace) | |
| assert workspace.id is not None | |
| assert workspace.name == "Test Workspace" | |
| assert workspace.status == WorkspaceStatus.ACTIVE.value | |
| assert workspace.plan_tier == "standard" | |
| assert workspace.is_startup is False | |
| assert workspace.learning_phase_completed is False | |
| assert workspace.metadata_json == {} | |
| assert workspace.created_at is not None | |
| # Note: updated_at is None on creation, only set on update | |
| def test_workspace_create_with_all_fields(self, db_session: Session): | |
| """Test Workspace creation with all optional fields.""" | |
| workspace = Workspace( | |
| name="Full Workspace", | |
| description="A workspace with all fields", | |
| status=WorkspaceStatus.ACTIVE.value, | |
| plan_tier="enterprise", | |
| is_startup=True, | |
| learning_phase_completed=True, | |
| metadata_json={"theme": "dark", "language": "en"}, | |
| internal_domains='["atom.ai", "example.com"]', | |
| ) | |
| db_session.add(workspace) | |
| db_session.commit() | |
| db_session.refresh(workspace) | |
| assert workspace.name == "Full Workspace" | |
| assert workspace.description == "A workspace with all fields" | |
| assert workspace.plan_tier == "enterprise" | |
| assert workspace.is_startup is True | |
| assert workspace.learning_phase_completed is True | |
| assert workspace.metadata_json["theme"] == "dark" | |
| assert workspace.internal_domains == '["atom.ai", "example.com"]' | |
| def test_workspace_tenant_relationship(self, db_session: Session): | |
| """Test Workspace belongs to Tenant (via tenant_id).""" | |
| tenant = TenantFactory(_session=db_session) | |
| db_session.commit() | |
| workspace = Workspace( | |
| name="Tenant Workspace", | |
| tenant_id=tenant.id, | |
| ) | |
| db_session.add(workspace) | |
| db_session.commit() | |
| db_session.refresh(workspace) | |
| assert workspace.tenant_id == tenant.id | |
| def test_workspace_teams_relationship(self, db_session: Session): | |
| """Test Workspace has many Teams (one-to-many).""" | |
| workspace = WorkspaceFactory(_session=db_session) | |
| db_session.commit() | |
| team1 = TeamFactory(workspace_id=workspace.id, _session=db_session) | |
| team2 = TeamFactory(workspace_id=workspace.id, _session=db_session) | |
| db_session.commit() | |
| retrieved_workspace = db_session.query(Workspace).filter( | |
| Workspace.id == workspace.id | |
| ).first() | |
| assert len(retrieved_workspace.teams) == 2 | |
| team_ids = [t.id for t in retrieved_workspace.teams] | |
| assert team1.id in team_ids | |
| assert team2.id in team_ids | |
| def test_workspace_users_relationship(self, db_session: Session): | |
| """Test Workspace has many Users (many-to-many via user_workspaces).""" | |
| workspace = WorkspaceFactory(_session=db_session) | |
| db_session.commit() | |
| user1 = UserFactory(email="user1@test.com", _session=db_session) | |
| user2 = UserFactory(email="user2@test.com", _session=db_session) | |
| workspace.users.append(user1) | |
| workspace.users.append(user2) | |
| db_session.commit() | |
| retrieved_workspace = db_session.query(Workspace).filter( | |
| Workspace.id == workspace.id | |
| ).first() | |
| assert len(retrieved_workspace.users) == 2 | |
| def test_workspace_is_startup_property(self, db_session: Session): | |
| """Test Workspace is_startup boolean property.""" | |
| workspace1 = Workspace(name="Startup Workspace", is_startup=True) | |
| workspace2 = Workspace(name="Corporate Workspace", is_startup=False) | |
| db_session.add(workspace1) | |
| db_session.add(workspace2) | |
| db_session.commit() | |
| assert workspace1.is_startup is True | |
| assert workspace2.is_startup is False | |
| def test_workspace_metadata_json_field(self, db_session: Session): | |
| """Test Workspace metadata_json field serialization.""" | |
| metadata = { | |
| "theme": "dark", | |
| "language": "en", | |
| "notifications": {"email": True, "push": False} | |
| } | |
| workspace = Workspace( | |
| name="Metadata Workspace", | |
| metadata_json=metadata, | |
| ) | |
| db_session.add(workspace) | |
| db_session.commit() | |
| db_session.refresh(workspace) | |
| assert isinstance(workspace.metadata_json, dict) | |
| assert workspace.metadata_json["theme"] == "dark" | |
| assert workspace.metadata_json["notifications"]["email"] is True | |
| class TestTeamModels: | |
| """Test Team model CRUD, relationships, and constraints.""" | |
| def test_team_create_with_defaults(self, db_session: Session): | |
| """Test Team creation with workspace.""" | |
| workspace = WorkspaceFactory(_session=db_session) | |
| db_session.commit() | |
| team = Team( | |
| name="Engineering Team", | |
| workspace_id=workspace.id, | |
| ) | |
| db_session.add(team) | |
| db_session.commit() | |
| db_session.refresh(team) | |
| assert team.id is not None | |
| assert team.name == "Engineering Team" | |
| assert team.workspace_id == workspace.id | |
| assert team.created_at is not None | |
| # Note: updated_at is None on creation, only set on update | |
| def test_team_members_relationship(self, db_session: Session): | |
| """Test Team has many Users (many-to-many via team_members).""" | |
| workspace = WorkspaceFactory(_session=db_session) | |
| team = TeamFactory(workspace_id=workspace.id, _session=db_session) | |
| db_session.commit() | |
| user1 = UserFactory(email="member1@test.com", _session=db_session) | |
| user2 = UserFactory(email="member2@test.com", _session=db_session) | |
| team.members.append(user1) | |
| team.members.append(user2) | |
| db_session.commit() | |
| retrieved_team = db_session.query(Team).filter(Team.id == team.id).first() | |
| assert len(retrieved_team.members) == 2 | |
| # Verify reverse relationship | |
| retrieved_user1 = db_session.query(User).filter(User.id == user1.id).first() | |
| assert len(retrieved_user1.teams) == 1 | |
| assert retrieved_user1.teams[0].id == team.id | |
| def test_team_workspace_relationship(self, db_session: Session): | |
| """Test Team belongs to Workspace (many-to-one).""" | |
| workspace = WorkspaceFactory(_session=db_session) | |
| team = TeamFactory(workspace_id=workspace.id, _session=db_session) | |
| db_session.commit() | |
| retrieved_team = db_session.query(Team).filter(Team.id == team.id).first() | |
| assert retrieved_team.workspace.id == workspace.id | |
| assert retrieved_team.workspace.name == workspace.name | |
| # ============================================================================ | |
| # Task 2: Tenant Models Tests | |
| # ============================================================================ | |
| class TestTenantModels: | |
| """Test Tenant model CRUD, properties, and relationships.""" | |
| def test_tenant_create_with_defaults(self, db_session: Session): | |
| """Test Tenant creation with required fields.""" | |
| tenant = Tenant( | |
| name="Test Tenant", | |
| subdomain="testtenant", | |
| ) | |
| db_session.add(tenant) | |
| db_session.commit() | |
| db_session.refresh(tenant) | |
| assert tenant.id is not None | |
| assert tenant.name == "Test Tenant" | |
| assert tenant.subdomain == "testtenant" | |
| assert tenant.plan_type == PlanType.FREE.value | |
| assert tenant.edition == "personal" | |
| assert tenant.memory_limit_mb == 50 | |
| assert tenant.is_active is True | |
| assert tenant.created_at is not None | |
| def test_tenant_edition_properties(self, db_session: Session): | |
| """Test Tenant edition properties (is_personal, is_enterprise).""" | |
| personal_tenant = Tenant( | |
| name="Personal Tenant", | |
| subdomain="personal", | |
| edition="personal", | |
| ) | |
| enterprise_tenant = Tenant( | |
| name="Enterprise Tenant", | |
| subdomain="enterprise", | |
| edition="enterprise", | |
| ) | |
| db_session.add(personal_tenant) | |
| db_session.add(enterprise_tenant) | |
| db_session.commit() | |
| assert personal_tenant.is_personal is True | |
| assert personal_tenant.is_enterprise is False | |
| assert enterprise_tenant.is_personal is False | |
| assert enterprise_tenant.is_enterprise is True | |
| def test_tenant_edition_display_name(self, db_session: Session): | |
| """Test Tenant edition_display_name property.""" | |
| personal_tenant = Tenant( | |
| name="Personal", | |
| subdomain="personal", | |
| edition="personal", | |
| ) | |
| enterprise_tenant = Tenant( | |
| name="Enterprise", | |
| subdomain="enterprise", | |
| edition="enterprise", | |
| ) | |
| db_session.add(personal_tenant) | |
| db_session.add(enterprise_tenant) | |
| db_session.commit() | |
| assert personal_tenant.edition_display_name == "Personal Edition" | |
| assert enterprise_tenant.edition_display_name == "Enterprise Edition" | |
| def test_tenant_can_upgrade_method(self, db_session: Session): | |
| """Test Tenant can_upgrade_to_enterprise() method.""" | |
| personal_tenant = Tenant( | |
| name="Personal", | |
| subdomain="personal", | |
| edition="personal", | |
| ) | |
| enterprise_tenant = Tenant( | |
| name="Enterprise", | |
| subdomain="enterprise", | |
| edition="enterprise", | |
| ) | |
| db_session.add(personal_tenant) | |
| db_session.add(enterprise_tenant) | |
| db_session.commit() | |
| assert personal_tenant.can_upgrade_to_enterprise() is True | |
| assert enterprise_tenant.can_upgrade_to_enterprise() is False | |
| def test_tenant_budget_fields(self, db_session: Session): | |
| """Test Tenant budget tracking fields.""" | |
| tenant = Tenant( | |
| name="Budget Tenant", | |
| subdomain="budget", | |
| budget_limit_usd=500.0, | |
| current_spend_usd=150.0, | |
| total_spend_usd=1000.0, | |
| ) | |
| db_session.add(tenant) | |
| db_session.commit() | |
| db_session.refresh(tenant) | |
| assert tenant.budget_limit_usd == 500.0 | |
| assert tenant.current_spend_usd == 150.0 | |
| assert tenant.total_spend_usd == 1000.0 | |
| # ============================================================================ | |
| # Task 2: UserAccount Models Tests | |
| # ============================================================================ | |
| class TestUserAccountModels: | |
| """Test UserAccount model for IM platform linking.""" | |
| def test_user_account_create(self, db_session: Session): | |
| """Test UserAccount creation for IM platform linking.""" | |
| user = UserFactory(_session=db_session) | |
| tenant = TenantFactory(_session=db_session) | |
| db_session.commit() | |
| user_account = UserAccount( | |
| user_id=user.id, | |
| tenant_id=tenant.id, | |
| platform="slack", | |
| platform_user_id="U12345", | |
| chat_id="C12345", | |
| username="testuser", | |
| ) | |
| db_session.add(user_account) | |
| db_session.commit() | |
| db_session.refresh(user_account) | |
| assert user_account.id is not None | |
| assert user_account.user_id == user.id | |
| assert user_account.tenant_id == tenant.id | |
| assert user_account.platform == "slack" | |
| assert user_account.platform_user_id == "U12345" | |
| assert user_account.chat_id == "C12345" | |
| assert user_account.username == "testuser" | |
| assert user_account.is_active is True | |
| assert user_account.linked_at is not None | |
| def test_user_account_unique_constraint(self, db_session: Session): | |
| """Test UserAccount unique constraint on (platform, platform_user_id).""" | |
| user = UserFactory(_session=db_session) | |
| tenant = TenantFactory(_session=db_session) | |
| db_session.commit() | |
| # Create first UserAccount | |
| user_account1 = UserAccount( | |
| user_id=user.id, | |
| tenant_id=tenant.id, | |
| platform="slack", | |
| platform_user_id="U12345", | |
| ) | |
| db_session.add(user_account1) | |
| db_session.commit() | |
| # Try to create duplicate (same platform + platform_user_id) | |
| with pytest.raises(IntegrityError): | |
| user_account2 = UserAccount( | |
| user_id=user.id, | |
| tenant_id=tenant.id, | |
| platform="slack", | |
| platform_user_id="U12345", # Same platform_user_id | |
| ) | |
| db_session.add(user_account2) | |
| db_session.commit() | |
| db_session.rollback() | |
| def test_user_account_user_relationship(self, db_session: Session): | |
| """Test UserAccount belongs to User.""" | |
| user = UserFactory(_session=db_session) | |
| tenant = TenantFactory(_session=db_session) | |
| db_session.commit() | |
| user_account = UserAccount( | |
| user_id=user.id, | |
| tenant_id=tenant.id, | |
| platform="discord", | |
| ) | |
| db_session.add(user_account) | |
| db_session.commit() | |
| assert user_account.user.id == user.id | |
| assert user_account.user.email == user.email | |
| def test_user_account_tenant_relationship(self, db_session: Session): | |
| """Test UserAccount belongs to Tenant.""" | |
| user = UserFactory(_session=db_session) | |
| tenant = TenantFactory(_session=db_session) | |
| db_session.commit() | |
| user_account = UserAccount( | |
| user_id=user.id, | |
| tenant_id=tenant.id, | |
| platform="teams", | |
| ) | |
| db_session.add(user_account) | |
| db_session.commit() | |
| assert user_account.tenant.id == tenant.id | |
| assert user_account.tenant.name == tenant.name | |
| # ============================================================================ | |
| # Task 2: OAuthToken Models Tests | |
| # ============================================================================ | |
| class TestOAuthTokenModels: | |
| """Test OAuthToken model for OAuth token management.""" | |
| def test_oauth_token_create(self, db_session: Session): | |
| """Test OAuthToken creation with provider.""" | |
| user = UserFactory(_session=db_session) | |
| tenant = TenantFactory(_session=db_session) | |
| db_session.commit() | |
| oauth_token = OAuthToken( | |
| client_id="client_123", | |
| user_id=user.id, | |
| tenant_id=tenant.id, | |
| access_token_hash="hash_" + "a" * 58, | |
| scope="read write", | |
| token_type="Bearer", | |
| access_token_expires_at=datetime.utcnow() + timedelta(hours=1), | |
| ) | |
| db_session.add(oauth_token) | |
| db_session.commit() | |
| db_session.refresh(oauth_token) | |
| assert oauth_token.id is not None | |
| assert oauth_token.user_id == user.id | |
| assert oauth_token.tenant_id == tenant.id | |
| assert oauth_token.access_token_hash.startswith("hash_") | |
| assert oauth_token.scope == "read write" | |
| assert oauth_token.token_type == "Bearer" | |
| assert oauth_token.access_token_expires_at is not None | |
| def test_oauth_token_user_relationship(self, db_session: Session): | |
| """Test OAuthToken belongs to User.""" | |
| user = UserFactory(_session=db_session) | |
| tenant = TenantFactory(_session=db_session) | |
| db_session.commit() | |
| oauth_token = OAuthToken( | |
| client_id="client_123", | |
| user_id=user.id, | |
| tenant_id=tenant.id, | |
| access_token_hash="hash_" + "a" * 58, | |
| scope="read", | |
| token_type="Bearer", | |
| access_token_expires_at=datetime.utcnow() + timedelta(hours=1), | |
| ) | |
| db_session.add(oauth_token) | |
| db_session.commit() | |
| # Note: OAuthToken model may not have a 'user' relationship defined | |
| # This test verifies the foreign key relationship works | |
| assert oauth_token.user_id == user.id | |
| def test_oauth_token_expires_at_field(self, db_session: Session): | |
| """Test OAuthToken datetime handling for expiration.""" | |
| user = UserFactory(_session=db_session) | |
| tenant = TenantFactory(_session=db_session) | |
| expires_at = datetime.utcnow() + timedelta(hours=2) | |
| db_session.commit() | |
| oauth_token = OAuthToken( | |
| client_id="client_123", | |
| user_id=user.id, | |
| tenant_id=tenant.id, | |
| access_token_hash="hash_" + "a" * 58, | |
| scope="read", | |
| token_type="Bearer", | |
| access_token_expires_at=expires_at, | |
| ) | |
| db_session.add(oauth_token) | |
| db_session.commit() | |
| db_session.refresh(oauth_token) | |
| assert oauth_token.access_token_expires_at is not None | |
| # Allow for small time differences during test execution | |
| time_diff = abs((oauth_token.access_token_expires_at - expires_at).total_seconds()) | |
| assert time_diff < 5 # Less than 5 seconds difference | |
| # ============================================================================ | |
| # Task 2: Chat Models Tests | |
| # ============================================================================ | |
| class TestChatModels: | |
| """Test ChatSession and ChatMessage models.""" | |
| def test_chat_session_create(self, db_session: Session): | |
| """Test ChatSession creation.""" | |
| chat_session = ChatSessionFactory(_session=db_session) | |
| db_session.commit() | |
| db_session.refresh(chat_session) | |
| assert chat_session.id is not None | |
| assert chat_session.user_id is not None | |
| assert chat_session.title is not None | |
| assert chat_session.created_at is not None | |
| assert chat_session.updated_at is not None | |
| assert chat_session.message_count >= 0 | |
| def test_chat_session_anonymous(self, db_session: Session): | |
| """Test ChatSession can be created without user (anonymous).""" | |
| # Note: ChatSession model requires user_id (nullable=False in schema) | |
| # This test verifies the model constraint | |
| with pytest.raises(IntegrityError): | |
| chat_session = ChatSession( | |
| id=None, # Required field | |
| user_id=None, # Should fail based on model schema | |
| ) | |
| db_session.add(chat_session) | |
| db_session.commit() | |
| db_session.rollback() | |
| def test_chat_session_messages_relationship(self, db_session: Session): | |
| """Test ChatSession has many Messages (one-to-many).""" | |
| chat_session = ChatSessionFactory(_session=db_session) | |
| db_session.commit() | |
| db_session.refresh(chat_session) | |
| # Note: ChatMessage model uses conversation_id, not session_id | |
| # This test verifies the relationship structure | |
| message1 = ChatMessageFactory( | |
| conversation_id=chat_session.id, | |
| role="user", | |
| content="Hello", | |
| _session=db_session, | |
| ) | |
| message2 = ChatMessageFactory( | |
| conversation_id=chat_session.id, | |
| role="assistant", | |
| content="Hi there!", | |
| _session=db_session, | |
| ) | |
| db_session.commit() | |
| # Verify messages were created with the conversation_id | |
| messages = db_session.query(ChatMessage).filter( | |
| ChatMessage.conversation_id == chat_session.id | |
| ).all() | |
| assert len(messages) == 2 | |
| def test_chat_message_create(self, db_session: Session): | |
| """Test ChatMessage creation with role and content.""" | |
| message = ChatMessageFactory( | |
| role="user", | |
| content="Test message", | |
| _session=db_session, | |
| ) | |
| db_session.commit() | |
| db_session.refresh(message) | |
| assert message.id is not None | |
| assert message.role == "user" | |
| assert message.content == "Test message" | |
| assert message.created_at is not None | |
| assert message.conversation_id is not None | |
| assert message.tenant_id is not None | |
| def test_chat_message_conversation_relationship(self, db_session: Session): | |
| """Test ChatMessage belongs to conversation (via conversation_id).""" | |
| chat_session = ChatSessionFactory(_session=db_session) | |
| db_session.commit() | |
| message = ChatMessageFactory( | |
| conversation_id=chat_session.id, | |
| role="assistant", | |
| content="Response", | |
| _session=db_session, | |
| ) | |
| db_session.commit() | |
| assert message.conversation_id == chat_session.id | |
| def test_chat_message_metadata_json(self, db_session: Session): | |
| """Test ChatMessage metadata_json field handling.""" | |
| metadata = '{"tokens": 150, "model": "gpt-4", "latency_ms": 250}' | |
| message = ChatMessage( | |
| conversation_id="conversation_123", | |
| tenant_id="tenant_123", | |
| role="assistant", | |
| content="Test", | |
| metadata_json=metadata, # Stored as Text | |
| ) | |
| db_session.add(message) | |
| db_session.commit() | |
| db_session.refresh(message) | |
| assert message.metadata_json is not None | |
| # Note: metadata_json is stored as Text, not JSON type | |
| assert isinstance(message.metadata_json, str) | |
| assert '"tokens": 150' in message.metadata_json | |