Refactor account models to use SQLAlchemy 2.0 dataclass mapping (#26415)

Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

api/tests/test_containers_integration_tests/services/test_account_service.py

@@ -16,6 +16,7 @@ from services.errors.account import (
     AccountPasswordError,
     AccountRegisterError,
     CurrentPasswordIncorrectError,
+    TenantNotFoundError,
 )
 from services.errors.workspace import WorkSpaceNotAllowedCreateError, WorkspacesLimitExceededError
 
@@ -1414,7 +1415,7 @@ class TestTenantService:
         )
 
         # Try to get current tenant (should fail)
-        with pytest.raises(AttributeError):
+        with pytest.raises((AttributeError, TenantNotFoundError)):
             TenantService.get_current_tenant_by_account(account)
 
     def test_switch_tenant_success(self, db_session_with_containers, mock_external_service_dependencies):

api/tests/test_containers_integration_tests/services/test_workflow_service.py

@@ -44,27 +44,26 @@ class TestWorkflowService:
             Account: Created test account instance
         """
         fake = fake or Faker()
-        account = Account()
-        account.id = fake.uuid4()
-        account.email = fake.email()
-        account.name = fake.name()
-        account.avatar_url = fake.url()
-        account.tenant_id = fake.uuid4()
-        account.status = "active"
-        account.type = "normal"
-        account.role = "owner"
-        account.interface_language = "en-US"  # Set interface language for Site creation
+        account = Account(
+            email=fake.email(),
+            name=fake.name(),
+            avatar=fake.url(),
+            status="active",
+            interface_language="en-US",  # Set interface language for Site creation
+        )
         account.created_at = fake.date_time_this_year()
+        account.id = fake.uuid4()
         account.updated_at = account.created_at
 
         # Create a tenant for the account
         from models.account import Tenant
 
-        tenant = Tenant()
-        tenant.id = account.tenant_id
-        tenant.name = f"Test Tenant {fake.company()}"
-        tenant.plan = "basic"
-        tenant.status = "active"
+        tenant = Tenant(
+            name=f"Test Tenant {fake.company()}",
+            plan="basic",
+            status="active",
+        )
+        tenant.id = account.current_tenant_id
         tenant.created_at = fake.date_time_this_year()
         tenant.updated_at = tenant.created_at
 
@@ -91,20 +90,21 @@ class TestWorkflowService:
             App: Created test app instance
         """
         fake = fake or Faker()
-        app = App()
-        app.id = fake.uuid4()
-        app.tenant_id = fake.uuid4()
-        app.name = fake.company()
-        app.description = fake.text()
-        app.mode = AppMode.WORKFLOW
-        app.icon_type = "emoji"
-        app.icon = "🤖"
-        app.icon_background = "#FFEAD5"
-        app.enable_site = True
-        app.enable_api = True
-        app.created_by = fake.uuid4()
+        app = App(
+            id=fake.uuid4(),
+            tenant_id=fake.uuid4(),
+            name=fake.company(),
+            description=fake.text(),
+            mode=AppMode.WORKFLOW,
+            icon_type="emoji",
+            icon="🤖",
+            icon_background="#FFEAD5",
+            enable_site=True,
+            enable_api=True,
+            created_by=fake.uuid4(),
+            workflow_id=None,  # Will be set when workflow is created
+        )
         app.updated_by = app.created_by
-        app.workflow_id = None  # Will be set when workflow is created
 
         from extensions.ext_database import db
 
@@ -126,19 +126,20 @@ class TestWorkflowService:
             Workflow: Created test workflow instance
         """
         fake = fake or Faker()
-        workflow = Workflow()
-        workflow.id = fake.uuid4()
-        workflow.tenant_id = app.tenant_id
-        workflow.app_id = app.id
-        workflow.type = WorkflowType.WORKFLOW.value
-        workflow.version = Workflow.VERSION_DRAFT
-        workflow.graph = json.dumps({"nodes": [], "edges": []})
-        workflow.features = json.dumps({"features": []})
-        # unique_hash is a computed property based on graph and features
-        workflow.created_by = account.id
-        workflow.updated_by = account.id
-        workflow.environment_variables = []
-        workflow.conversation_variables = []
+        workflow = Workflow(
+            id=fake.uuid4(),
+            tenant_id=app.tenant_id,
+            app_id=app.id,
+            type=WorkflowType.WORKFLOW.value,
+            version=Workflow.VERSION_DRAFT,
+            graph=json.dumps({"nodes": [], "edges": []}),
+            features=json.dumps({"features": []}),
+            # unique_hash is a computed property based on graph and features
+            created_by=account.id,
+            updated_by=account.id,
+            environment_variables=[],
+            conversation_variables=[],
+        )
 
         from extensions.ext_database import db
 

api/tests/test_containers_integration_tests/tasks/test_delete_segment_from_index_task.py

@@ -48,11 +48,8 @@ class TestDeleteSegmentFromIndexTask:
             Tenant: Created test tenant instance
         """
         fake = fake or Faker()
-        tenant = Tenant()
+        tenant = Tenant(name=f"Test Tenant {fake.company()}", plan="basic", status="active")
         tenant.id = fake.uuid4()
-        tenant.name = f"Test Tenant {fake.company()}"
-        tenant.plan = "basic"
-        tenant.status = "active"
         tenant.created_at = fake.date_time_this_year()
         tenant.updated_at = tenant.created_at
 
@@ -73,16 +70,14 @@ class TestDeleteSegmentFromIndexTask:
             Account: Created test account instance
         """
         fake = fake or Faker()
-        account = Account()
+        account = Account(
+            name=fake.name(),
+            email=fake.email(),
+            avatar=fake.url(),
+            status="active",
+            interface_language="en-US",
+        )
         account.id = fake.uuid4()
-        account.email = fake.email()
-        account.name = fake.name()
-        account.avatar_url = fake.url()
-        account.tenant_id = tenant.id
-        account.status = "active"
-        account.type = "normal"
-        account.role = "owner"
-        account.interface_language = "en-US"
         account.created_at = fake.date_time_this_year()
         account.updated_at = account.created_at
 

api/tests/test_containers_integration_tests/tasks/test_disable_segments_from_index_task.py

@@ -43,27 +43,30 @@ class TestDisableSegmentsFromIndexTask:
             Account: Created test account instance
         """
         fake = fake or Faker()
-        account = Account()
+        account = Account(
+            email=fake.email(),
+            name=fake.name(),
+            avatar=fake.url(),
+            status="active",
+            interface_language="en-US",
+        )
         account.id = fake.uuid4()
-        account.email = fake.email()
-        account.name = fake.name()
-        account.avatar_url = fake.url()
+        # monkey-patch attributes for test setup
         account.tenant_id = fake.uuid4()
-        account.status = "active"
         account.type = "normal"
         account.role = "owner"
-        account.interface_language = "en-US"
         account.created_at = fake.date_time_this_year()
         account.updated_at = account.created_at
 
         # Create a tenant for the account
         from models.account import Tenant
 
-        tenant = Tenant()
+        tenant = Tenant(
+            name=f"Test Tenant {fake.company()}",
+            plan="basic",
+            status="active",
+        )
         tenant.id = account.tenant_id
-        tenant.name = f"Test Tenant {fake.company()}"
-        tenant.plan = "basic"
-        tenant.status = "active"
         tenant.created_at = fake.date_time_this_year()
         tenant.updated_at = tenant.created_at
 
@@ -91,20 +94,21 @@ class TestDisableSegmentsFromIndexTask:
             Dataset: Created test dataset instance
         """
         fake = fake or Faker()
-        dataset = Dataset()
-        dataset.id = fake.uuid4()
-        dataset.tenant_id = account.tenant_id
-        dataset.name = f"Test Dataset {fake.word()}"
-        dataset.description = fake.text(max_nb_chars=200)
-        dataset.provider = "vendor"
-        dataset.permission = "only_me"
-        dataset.data_source_type = "upload_file"
-        dataset.indexing_technique = "high_quality"
-        dataset.created_by = account.id
-        dataset.updated_by = account.id
-        dataset.embedding_model = "text-embedding-ada-002"
-        dataset.embedding_model_provider = "openai"
-        dataset.built_in_field_enabled = False
+        dataset = Dataset(
+            id=fake.uuid4(),
+            tenant_id=account.tenant_id,
+            name=f"Test Dataset {fake.word()}",
+            description=fake.text(max_nb_chars=200),
+            provider="vendor",
+            permission="only_me",
+            data_source_type="upload_file",
+            indexing_technique="high_quality",
+            created_by=account.id,
+            updated_by=account.id,
+            embedding_model="text-embedding-ada-002",
+            embedding_model_provider="openai",
+            built_in_field_enabled=False,
+        )
 
         from extensions.ext_database import db
 
@@ -128,6 +132,7 @@ class TestDisableSegmentsFromIndexTask:
         """
         fake = fake or Faker()
         document = DatasetDocument()
+
         document.id = fake.uuid4()
         document.tenant_id = dataset.tenant_id
         document.dataset_id = dataset.id
@@ -153,7 +158,6 @@ class TestDisableSegmentsFromIndexTask:
         document.archived = False
         document.doc_form = "text_model"  # Use text_model form for testing
         document.doc_language = "en"
-
         from extensions.ext_database import db
 
         db.session.add(document)

api/tests/test_containers_integration_tests/tasks/test_mail_invite_member_task.py

@@ -96,9 +96,9 @@ class TestMailInviteMemberTask:
             password=fake.password(),
             interface_language="en-US",
             status=AccountStatus.ACTIVE.value,
-            created_at=datetime.now(UTC),
-            updated_at=datetime.now(UTC),
         )
+        account.created_at = datetime.now(UTC)
+        account.updated_at = datetime.now(UTC)
         db_session_with_containers.add(account)
         db_session_with_containers.commit()
         db_session_with_containers.refresh(account)
@@ -106,9 +106,9 @@ class TestMailInviteMemberTask:
         # Create tenant
         tenant = Tenant(
             name=fake.company(),
-            created_at=datetime.now(UTC),
-            updated_at=datetime.now(UTC),
         )
+        tenant.created_at = datetime.now(UTC)
+        tenant.updated_at = datetime.now(UTC)
         db_session_with_containers.add(tenant)
         db_session_with_containers.commit()
         db_session_with_containers.refresh(tenant)
@@ -118,8 +118,8 @@ class TestMailInviteMemberTask:
             tenant_id=tenant.id,
             account_id=account.id,
             role=TenantAccountRole.OWNER.value,
-            created_at=datetime.now(UTC),
         )
+        tenant_join.created_at = datetime.now(UTC)
         db_session_with_containers.add(tenant_join)
         db_session_with_containers.commit()
 
@@ -164,9 +164,10 @@ class TestMailInviteMemberTask:
             password="",
             interface_language="en-US",
             status=AccountStatus.PENDING.value,
-            created_at=datetime.now(UTC),
-            updated_at=datetime.now(UTC),
         )
+
+        account.created_at = datetime.now(UTC)
+        account.updated_at = datetime.now(UTC)
         db_session_with_containers.add(account)
         db_session_with_containers.commit()
         db_session_with_containers.refresh(account)
@@ -176,8 +177,8 @@ class TestMailInviteMemberTask:
             tenant_id=tenant.id,
             account_id=account.id,
             role=TenantAccountRole.NORMAL.value,
-            created_at=datetime.now(UTC),
         )
+        tenant_join.created_at = datetime.now(UTC)
         db_session_with_containers.add(tenant_join)
         db_session_with_containers.commit()