From b58d9e030a6fc123845073eae6472763f9f6042e Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Fri, 30 Jan 2026 22:39:02 +0900 Subject: [PATCH] refactor: init_validate.py to v3 (#31457) --- api/controllers/console/init_validate.py | 93 ++++++++----------- api/extensions/ext_fastopenapi.py | 2 + .../console/test_fastopenapi_init_validate.py | 46 +++++++++ 3 files changed, 88 insertions(+), 53 deletions(-) create mode 100644 api/tests/unit_tests/controllers/console/test_fastopenapi_init_validate.py diff --git a/api/controllers/console/init_validate.py b/api/controllers/console/init_validate.py index 2bebe79eac..f086bf1862 100644 --- a/api/controllers/console/init_validate.py +++ b/api/controllers/console/init_validate.py @@ -1,87 +1,74 @@ import os +from typing import Literal from flask import session -from flask_restx import Resource, fields from pydantic import BaseModel, Field from sqlalchemy import select from sqlalchemy.orm import Session from configs import dify_config +from controllers.fastopenapi import console_router from extensions.ext_database import db from models.model import DifySetup from services.account_service import TenantService -from . import console_ns from .error import AlreadySetupError, InitValidateFailedError from .wraps import only_edition_self_hosted -DEFAULT_REF_TEMPLATE_SWAGGER_2_0 = "#/definitions/{model}" - class InitValidatePayload(BaseModel): - password: str = Field(..., max_length=30) + password: str = Field(..., max_length=30, description="Initialization password") -console_ns.schema_model( - InitValidatePayload.__name__, - InitValidatePayload.model_json_schema(ref_template=DEFAULT_REF_TEMPLATE_SWAGGER_2_0), +class InitStatusResponse(BaseModel): + status: Literal["finished", "not_started"] = Field(..., description="Initialization status") + + +class InitValidateResponse(BaseModel): + result: str = Field(description="Operation result", examples=["success"]) + + +@console_router.get( + "/init", + response_model=InitStatusResponse, + tags=["console"], ) +def get_init_status() -> InitStatusResponse: + """Get initialization validation status.""" + init_status = get_init_validate_status() + if init_status: + return InitStatusResponse(status="finished") + return InitStatusResponse(status="not_started") -@console_ns.route("/init") -class InitValidateAPI(Resource): - @console_ns.doc("get_init_status") - @console_ns.doc(description="Get initialization validation status") - @console_ns.response( - 200, - "Success", - model=console_ns.model( - "InitStatusResponse", - {"status": fields.String(description="Initialization status", enum=["finished", "not_started"])}, - ), - ) - def get(self): - """Get initialization validation status""" - init_status = get_init_validate_status() - if init_status: - return {"status": "finished"} - return {"status": "not_started"} +@console_router.post( + "/init", + response_model=InitValidateResponse, + tags=["console"], + status_code=201, +) +@only_edition_self_hosted +def validate_init_password(payload: InitValidatePayload) -> InitValidateResponse: + """Validate initialization password.""" + tenant_count = TenantService.get_tenant_count() + if tenant_count > 0: + raise AlreadySetupError() - @console_ns.doc("validate_init_password") - @console_ns.doc(description="Validate initialization password for self-hosted edition") - @console_ns.expect(console_ns.models[InitValidatePayload.__name__]) - @console_ns.response( - 201, - "Success", - model=console_ns.model("InitValidateResponse", {"result": fields.String(description="Operation result")}), - ) - @console_ns.response(400, "Already setup or validation failed") - @only_edition_self_hosted - def post(self): - """Validate initialization password""" - # is tenant created - tenant_count = TenantService.get_tenant_count() - if tenant_count > 0: - raise AlreadySetupError() + if payload.password != os.environ.get("INIT_PASSWORD"): + session["is_init_validated"] = False + raise InitValidateFailedError() - payload = InitValidatePayload.model_validate(console_ns.payload) - input_password = payload.password - - if input_password != os.environ.get("INIT_PASSWORD"): - session["is_init_validated"] = False - raise InitValidateFailedError() - - session["is_init_validated"] = True - return {"result": "success"}, 201 + session["is_init_validated"] = True + return InitValidateResponse(result="success") -def get_init_validate_status(): +def get_init_validate_status() -> bool: if dify_config.EDITION == "SELF_HOSTED": if os.environ.get("INIT_PASSWORD"): if session.get("is_init_validated"): return True with Session(db.engine) as db_session: - return db_session.execute(select(DifySetup)).scalar_one_or_none() + return db_session.execute(select(DifySetup)).scalar_one_or_none() is not None return True diff --git a/api/extensions/ext_fastopenapi.py b/api/extensions/ext_fastopenapi.py index 719456803a..ab4d23a072 100644 --- a/api/extensions/ext_fastopenapi.py +++ b/api/extensions/ext_fastopenapi.py @@ -27,9 +27,11 @@ def init_app(app: DifyApp) -> None: ) # Ensure route decorators are evaluated. + import controllers.console.init_validate as init_validate_module import controllers.console.ping as ping_module from controllers.console import remote_files, setup + _ = init_validate_module _ = ping_module _ = remote_files _ = setup diff --git a/api/tests/unit_tests/controllers/console/test_fastopenapi_init_validate.py b/api/tests/unit_tests/controllers/console/test_fastopenapi_init_validate.py new file mode 100644 index 0000000000..b9bc42fb25 --- /dev/null +++ b/api/tests/unit_tests/controllers/console/test_fastopenapi_init_validate.py @@ -0,0 +1,46 @@ +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 + app.secret_key = "test-secret-key" + return app + + +def test_console_init_get_returns_finished_when_no_init_password(app: Flask, monkeypatch: pytest.MonkeyPatch): + ext_fastopenapi.init_app(app) + monkeypatch.delenv("INIT_PASSWORD", raising=False) + + with patch("controllers.console.init_validate.dify_config.EDITION", "SELF_HOSTED"): + client = app.test_client() + response = client.get("/console/api/init") + + assert response.status_code == 200 + assert response.get_json() == {"status": "finished"} + + +def test_console_init_post_returns_success(app: Flask, monkeypatch: pytest.MonkeyPatch): + ext_fastopenapi.init_app(app) + monkeypatch.setenv("INIT_PASSWORD", "test-init-password") + + with ( + patch("controllers.console.init_validate.dify_config.EDITION", "SELF_HOSTED"), + patch("controllers.console.init_validate.TenantService.get_tenant_count", return_value=0), + ): + client = app.test_client() + response = client.post("/console/api/init", json={"password": "test-init-password"}) + + assert response.status_code == 201 + assert response.get_json() == {"result": "success"}