You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
116 lines
3.4 KiB
Python
116 lines
3.4 KiB
Python
import pathlib
|
|
import random
|
|
import secrets
|
|
from collections.abc import Generator
|
|
|
|
import pytest
|
|
from flask import Flask
|
|
from flask.testing import FlaskClient
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app_factory import create_app
|
|
from models import Account, DifySetup, Tenant, TenantAccountJoin, db
|
|
from services.account_service import AccountService, RegisterService
|
|
|
|
|
|
# Loading the .env file if it exists
|
|
def _load_env() -> None:
|
|
current_file_path = pathlib.Path(__file__).absolute()
|
|
# Items later in the list have higher precedence.
|
|
files_to_load = [".env", "vdb.env"]
|
|
|
|
env_file_paths = [current_file_path.parent / i for i in files_to_load]
|
|
for path in env_file_paths:
|
|
if not path.exists():
|
|
continue
|
|
|
|
from dotenv import load_dotenv
|
|
|
|
# Set `override=True` to ensure values from `vdb.env` take priority over values from `.env`
|
|
load_dotenv(str(path), override=True)
|
|
|
|
|
|
_load_env()
|
|
|
|
_CACHED_APP = create_app()
|
|
|
|
|
|
@pytest.fixture
|
|
def flask_app() -> Flask:
|
|
return _CACHED_APP
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def setup_account(request) -> Generator[Account, None, None]:
|
|
"""`dify_setup` completes the setup process for the Dify application.
|
|
|
|
It creates `Account` and `Tenant`, and inserts a `DifySetup` record into the database.
|
|
|
|
Most tests in the `controllers` package may require dify has been successfully setup.
|
|
"""
|
|
with _CACHED_APP.test_request_context():
|
|
# Check if setup already exists
|
|
existing_setup = db.session.query(DifySetup).first()
|
|
if existing_setup:
|
|
# Get the first available account
|
|
account = db.session.query(Account).first()
|
|
if account:
|
|
yield account
|
|
return
|
|
|
|
rand_suffix = random.randint(int(1e6), int(1e7)) # noqa
|
|
name = f"test-user-{rand_suffix}"
|
|
email = f"{name}@example.com"
|
|
|
|
# Clean up any existing setup first
|
|
from models.account import AccountMFASettings
|
|
|
|
db.session.query(AccountMFASettings).delete()
|
|
db.session.query(DifySetup).delete()
|
|
db.session.query(TenantAccountJoin).delete()
|
|
db.session.query(Account).delete()
|
|
db.session.query(Tenant).delete()
|
|
db.session.commit()
|
|
|
|
RegisterService.setup(
|
|
email=email,
|
|
name=name,
|
|
password=secrets.token_hex(16),
|
|
ip_address="localhost",
|
|
)
|
|
|
|
with _CACHED_APP.test_request_context():
|
|
with Session(bind=db.engine, expire_on_commit=False) as session:
|
|
account = session.query(Account).filter_by(email=email).one()
|
|
|
|
yield account
|
|
|
|
with _CACHED_APP.test_request_context():
|
|
# Clean up MFA settings first to avoid foreign key violations
|
|
from models.account import AccountMFASettings
|
|
|
|
db.session.query(AccountMFASettings).delete()
|
|
db.session.query(DifySetup).delete()
|
|
db.session.query(TenantAccountJoin).delete()
|
|
db.session.query(Account).delete()
|
|
db.session.query(Tenant).delete()
|
|
db.session.commit()
|
|
|
|
|
|
@pytest.fixture
|
|
def flask_req_ctx():
|
|
with _CACHED_APP.test_request_context():
|
|
yield
|
|
|
|
|
|
@pytest.fixture
|
|
def auth_header(setup_account) -> dict[str, str]:
|
|
token = AccountService.get_account_jwt_token(setup_account)
|
|
return {"Authorization": f"Bearer {token}"}
|
|
|
|
|
|
@pytest.fixture
|
|
def test_client() -> Generator[FlaskClient, None, None]:
|
|
with _CACHED_APP.test_client() as client:
|
|
yield client
|