diff --git a/api/controllers/console/setup.py b/api/controllers/console/setup.py index ed22ef045d..e1ea007232 100644 --- a/api/controllers/console/setup.py +++ b/api/controllers/console/setup.py @@ -1,20 +1,19 @@ +from typing import Literal + from flask import request -from flask_restx import Resource, fields from pydantic import BaseModel, Field, field_validator from configs import dify_config +from controllers.fastopenapi import console_router from libs.helper import EmailStr, extract_remote_ip from libs.password import valid_password from models.model import DifySetup, db from services.account_service import RegisterService, TenantService -from . import console_ns from .error import AlreadySetupError, NotInitValidateError from .init_validate import get_init_validate_status from .wraps import only_edition_self_hosted -DEFAULT_REF_TEMPLATE_SWAGGER_2_0 = "#/definitions/{model}" - class SetupRequestPayload(BaseModel): email: EmailStr = Field(..., description="Admin email address") @@ -28,78 +27,66 @@ class SetupRequestPayload(BaseModel): return valid_password(value) -console_ns.schema_model( - SetupRequestPayload.__name__, - SetupRequestPayload.model_json_schema(ref_template=DEFAULT_REF_TEMPLATE_SWAGGER_2_0), +class SetupStatusResponse(BaseModel): + step: Literal["not_started", "finished"] = Field(description="Setup step status") + setup_at: str | None = Field(default=None, description="Setup completion time (ISO format)") + + +class SetupResponse(BaseModel): + result: str = Field(description="Setup result", examples=["success"]) + + +@console_router.get( + "/setup", + response_model=SetupStatusResponse, + tags=["console"], ) +def get_setup_status_api() -> SetupStatusResponse: + """Get system setup status.""" + if dify_config.EDITION == "SELF_HOSTED": + setup_status = get_setup_status() + if setup_status and not isinstance(setup_status, bool): + return SetupStatusResponse(step="finished", setup_at=setup_status.setup_at.isoformat()) + if setup_status: + return SetupStatusResponse(step="finished") + return SetupStatusResponse(step="not_started") + return SetupStatusResponse(step="finished") -@console_ns.route("/setup") -class SetupApi(Resource): - @console_ns.doc("get_setup_status") - @console_ns.doc(description="Get system setup status") - @console_ns.response( - 200, - "Success", - console_ns.model( - "SetupStatusResponse", - { - "step": fields.String(description="Setup step status", enum=["not_started", "finished"]), - "setup_at": fields.String(description="Setup completion time (ISO format)", required=False), - }, - ), +@console_router.post( + "/setup", + response_model=SetupResponse, + tags=["console"], + status_code=201, +) +@only_edition_self_hosted +def setup_system(payload: SetupRequestPayload) -> SetupResponse: + """Initialize system setup with admin account.""" + if get_setup_status(): + raise AlreadySetupError() + + tenant_count = TenantService.get_tenant_count() + if tenant_count > 0: + raise AlreadySetupError() + + if not get_init_validate_status(): + raise NotInitValidateError() + + normalized_email = payload.email.lower() + + RegisterService.setup( + email=normalized_email, + name=payload.name, + password=payload.password, + ip_address=extract_remote_ip(request), + language=payload.language, ) - def get(self): - """Get system setup status""" - if dify_config.EDITION == "SELF_HOSTED": - setup_status = get_setup_status() - # Check if setup_status is a DifySetup object rather than a bool - if setup_status and not isinstance(setup_status, bool): - return {"step": "finished", "setup_at": setup_status.setup_at.isoformat()} - elif setup_status: - return {"step": "finished"} - return {"step": "not_started"} - return {"step": "finished"} - @console_ns.doc("setup_system") - @console_ns.doc(description="Initialize system setup with admin account") - @console_ns.expect(console_ns.models[SetupRequestPayload.__name__]) - @console_ns.response( - 201, "Success", console_ns.model("SetupResponse", {"result": fields.String(description="Setup result")}) - ) - @console_ns.response(400, "Already setup or validation failed") - @only_edition_self_hosted - def post(self): - """Initialize system setup with admin account""" - # is set up - if get_setup_status(): - raise AlreadySetupError() - - # is tenant created - tenant_count = TenantService.get_tenant_count() - if tenant_count > 0: - raise AlreadySetupError() - - if not get_init_validate_status(): - raise NotInitValidateError() - - args = SetupRequestPayload.model_validate(console_ns.payload) - normalized_email = args.email.lower() - - # setup - RegisterService.setup( - email=normalized_email, - name=args.name, - password=args.password, - ip_address=extract_remote_ip(request), - language=args.language, - ) - - return {"result": "success"}, 201 + return SetupResponse(result="success") -def get_setup_status(): +def get_setup_status() -> DifySetup | bool | None: if dify_config.EDITION == "SELF_HOSTED": return db.session.query(DifySetup).first() - else: - return True + + return True diff --git a/api/extensions/ext_fastopenapi.py b/api/extensions/ext_fastopenapi.py index 0ef1513e11..5f98aa7b67 100644 --- a/api/extensions/ext_fastopenapi.py +++ b/api/extensions/ext_fastopenapi.py @@ -28,8 +28,10 @@ def init_app(app: DifyApp) -> None: # Ensure route decorators are evaluated. import controllers.console.ping as ping_module + from controllers.console import setup _ = ping_module + _ = setup router.include_router(console_router, prefix="/console/api") CORS( diff --git a/api/tests/unit_tests/controllers/console/test_fastopenapi_setup.py b/api/tests/unit_tests/controllers/console/test_fastopenapi_setup.py new file mode 100644 index 0000000000..385539b6f3 --- /dev/null +++ b/api/tests/unit_tests/controllers/console/test_fastopenapi_setup.py @@ -0,0 +1,56 @@ +import builtins +from unittest.mock import patch + +import pytest +from flask import Flask +from flask.views import MethodView + +from extensions import ext_fastopenapi + +if not hasattr(builtins, "MethodView"): + builtins.MethodView = MethodView # type: ignore[attr-defined] + + +@pytest.fixture +def app() -> Flask: + app = Flask(__name__) + app.config["TESTING"] = True + return app + + +def test_console_setup_fastopenapi_get_not_started(app: Flask): + ext_fastopenapi.init_app(app) + + with ( + patch("controllers.console.setup.dify_config.EDITION", "SELF_HOSTED"), + patch("controllers.console.setup.get_setup_status", return_value=None), + ): + client = app.test_client() + response = client.get("/console/api/setup") + + assert response.status_code == 200 + assert response.get_json() == {"step": "not_started", "setup_at": None} + + +def test_console_setup_fastopenapi_post_success(app: Flask): + ext_fastopenapi.init_app(app) + + payload = { + "email": "admin@example.com", + "name": "Admin", + "password": "Passw0rd1", + "language": "en-US", + } + + with ( + patch("controllers.console.wraps.dify_config.EDITION", "SELF_HOSTED"), + patch("controllers.console.setup.get_setup_status", return_value=None), + patch("controllers.console.setup.TenantService.get_tenant_count", return_value=0), + patch("controllers.console.setup.get_init_validate_status", return_value=True), + patch("controllers.console.setup.RegisterService.setup"), + ): + client = app.test_client() + response = client.post("/console/api/setup", json=payload) + + assert response.status_code == 201 + assert response.get_json() == {"result": "success"} diff --git a/api/tests/unit_tests/controllers/console/test_setup.py b/api/tests/unit_tests/controllers/console/test_setup.py deleted file mode 100644 index e7882dcd2b..0000000000 --- a/api/tests/unit_tests/controllers/console/test_setup.py +++ /dev/null @@ -1,39 +0,0 @@ -from types import SimpleNamespace -from unittest.mock import patch - -from controllers.console.setup import SetupApi - - -class TestSetupApi: - def test_post_lowercases_email_before_register(self): - """Ensure setup registration normalizes email casing.""" - payload = { - "email": "Admin@Example.com", - "name": "Admin User", - "password": "ValidPass123!", - "language": "en-US", - } - setup_api = SetupApi(api=None) - - mock_console_ns = SimpleNamespace(payload=payload) - - with ( - patch("controllers.console.setup.console_ns", mock_console_ns), - patch("controllers.console.setup.get_setup_status", return_value=False), - patch("controllers.console.setup.TenantService.get_tenant_count", return_value=0), - patch("controllers.console.setup.get_init_validate_status", return_value=True), - patch("controllers.console.setup.extract_remote_ip", return_value="127.0.0.1"), - patch("controllers.console.setup.request", object()), - patch("controllers.console.setup.RegisterService.setup") as mock_register, - ): - response, status = setup_api.post() - - assert response == {"result": "success"} - assert status == 201 - mock_register.assert_called_once_with( - email="admin@example.com", - name=payload["name"], - password=payload["password"], - ip_address="127.0.0.1", - language=payload["language"], - )