dify/api/tests/test_containers_integration_tests/controllers/console/helpers.py

"""Shared helpers for authenticated console controller integration tests."""

import uuid

from flask.testing import FlaskClient
from sqlalchemy import select
from sqlalchemy.orm import Session

from configs import dify_config
from constants import HEADER_NAME_CSRF_TOKEN
from libs.datetime_utils import naive_utc_now
from libs.token import _real_cookie_name, generate_csrf_token
from models import Account, DifySetup, Tenant, TenantAccountJoin
from models.account import AccountStatus, TenantAccountRole, TenantStatus
from models.model import App, AppMode
from services.account_service import AccountService


def ensure_dify_setup(db_session: Session) -> None:
    """Create a setup marker once so setup-protected console routes can be exercised."""
    if db_session.scalar(select(DifySetup).limit(1)) is not None:
        return

    db_session.add(DifySetup(version=dify_config.project.version))
    db_session.commit()


def create_console_account_and_tenant(db_session: Session) -> tuple[Account, Tenant]:
    """Create an initialized owner account with a current tenant."""
    account = Account(
        email=f"test-{uuid.uuid4()}@example.com",
        name="Test User",
        interface_language="en-US",
        status=AccountStatus.ACTIVE,
    )
    account.initialized_at = naive_utc_now()
    db_session.add(account)
    db_session.commit()

    tenant = Tenant(name="Test Tenant", status=TenantStatus.NORMAL)
    db_session.add(tenant)
    db_session.commit()

    db_session.add(
        TenantAccountJoin(
            tenant_id=tenant.id,
            account_id=account.id,
            role=TenantAccountRole.OWNER,
            current=True,
        )
    )
    db_session.commit()

    account.set_tenant_id(tenant.id)
    account.timezone = "UTC"
    db_session.commit()

    ensure_dify_setup(db_session)
    return account, tenant


def create_console_app(db_session: Session, tenant_id: str, account_id: str, mode: AppMode) -> App:
    """Create a minimal app row that can be loaded by get_app_model."""
    app = App(
        tenant_id=tenant_id,
        name="Test App",
        mode=mode,
        enable_site=True,
        enable_api=True,
        created_by=account_id,
    )
    db_session.add(app)
    db_session.commit()
    return app


def authenticate_console_client(test_client: FlaskClient, account: Account) -> dict[str, str]:
    """Attach console auth cookies/headers for endpoints guarded by login_required."""
    access_token = AccountService.get_account_jwt_token(account)
    csrf_token = generate_csrf_token(account.id)
    test_client.set_cookie(_real_cookie_name("csrf_token"), csrf_token, domain="localhost")
    return {
        "Authorization": f"Bearer {access_token}",
        HEADER_NAME_CSRF_TOKEN: csrf_token,
    }