dify/api/tests/test_containers_integration_tests/controllers/openapi/conftest.py

from __future__ import annotations

import uuid
from collections.abc import Callable
from unittest.mock import patch

import pytest
from faker import Faker
from flask import Flask
from sqlalchemy.orm import Session

from controllers.openapi.auth.context import Context
from controllers.openapi.auth.loaders import PathParam, load_app
from controllers.openapi.auth.requirements import ResolveCaller
from controllers.openapi.auth.subjects import subject_from_auth
from libs.oauth_bearer import AuthContext, TokenType
from models import Account, Tenant
from services.account_service import AccountService, TenantService
from tests.test_containers_integration_tests.helpers import generate_valid_password

_CLIENT_ID = "integration-cli"


@pytest.fixture
def app(flask_app_with_containers: Flask) -> Flask:
    return flask_app_with_containers


@pytest.fixture
def make_account(db_session_with_containers: Session) -> Callable[..., Account]:
    """Factory that registers a real Account and gives it an owner workspace.

    System feature gates are stubbed (registration / workspace creation
    allowed) exactly like the AppDslService integration tests, so this stays a
    pure account+tenant setup helper.
    """

    # Depend on db_session_with_containers so the app context / DB session is
    # active for the real AccountService/TenantService calls below.
    assert db_session_with_containers is not None

    def _make(*, with_owner_tenant: bool = True) -> Account:
        fake = Faker()
        with patch("services.account_service.SystemFeatureService") as mock_feature_service:
            mock_feature_service.is_registration_allowed.return_value = True
            account = AccountService.create_account(
                email=fake.email(),
                name=fake.name(),
                interface_language="en-US",
                password=generate_valid_password(fake),
                session=db_session_with_containers,
            )
            if with_owner_tenant:
                TenantService.create_owner_tenant_if_not_exist(
                    account, name=fake.company(), session=db_session_with_containers
                )
        return account

    return _make


def add_tenant_for_account(
    account: Account, *, session: Session, role: str = "normal", name: str = "Second WS"
) -> Tenant:
    """Create an additional tenant and join ``account`` to it (real service calls)."""
    with patch("services.account_service.SystemFeatureService") as mock_feature_service:
        mock_feature_service.is_workspace_creation_allowed.return_value = True
        tenant = TenantService.create_tenant(name=name, session=session)
    TenantService.create_tenant_member(tenant, account, session, role=role)
    return tenant


def _account_auth(
    account: Account,
    *,
    token_id: uuid.UUID | None = None,
    client_id: str = _CLIENT_ID,
) -> AuthContext:
    """The ``AuthContext`` a live ``dfoa_`` token for ``account`` would carry."""
    return AuthContext(
        subject_email=account.email,
        subject_issuer=None,
        account_id=uuid.UUID(str(account.id)),
        client_id=client_id,
        token_id=token_id or uuid.uuid4(),
        token_type=TokenType.OAUTH_ACCOUNT,
        expires_at=None,
    )


def context_for(
    account: Account,
    *,
    session: Session,
    view_args: dict[str, str] | None = None,
    token_id: uuid.UUID | None = None,
) -> Context:
    """Build the ``Context`` a handler is given after the pipeline ran.

    The subject comes from a real ``AuthContext`` through ``subject_from_auth``,
    so the helper walks the same resolution path the router does. ``view_args``
    is the route's path params — the app and the workspace are loaded from it,
    so a route carrying ``<app_id>`` needs ``{"app_id": ...}`` here.
    ``token_id`` only matters to the ``/account/sessions*`` family, which reads
    it back off the subject.

    It runs ``ResolveCaller`` itself — the requirement every pipeline fixes
    last — rather than the wider set a route's own requirements would ask for,
    so a handler sees exactly what the thinnest pipeline would give it. On a
    route carrying ``<app_id>`` that thinnest pipeline is ``CheckAppApiEnabled``,
    which loads the app; handlers read ``ctx.app`` and never load it themselves,
    so the helper loads it through the same loader. Running the real pieces is
    what keeps this CI-only helper from drifting away from the pipeline it
    stands in for.
    """
    ctx = Context(subject_from_auth(_account_auth(account, token_id=token_id)), session, view_args or {})
    if PathParam.APP_ID in ctx.view_args:
        load_app(ctx)
    ResolveCaller().run(ctx.subject, ctx, session)
    return ctx