mirror of
https://github.com/langgenius/dify.git
synced 2026-08-31 13:06:11 +08:00
251 lines
9.1 KiB
Python
251 lines
9.1 KiB
Python
from __future__ import annotations
|
|
|
|
import builtins
|
|
import uuid
|
|
from datetime import UTC, datetime
|
|
from unittest.mock import ANY, MagicMock
|
|
|
|
import pytest
|
|
from flask import Flask
|
|
from flask.views import MethodView as FlaskMethodView
|
|
|
|
from tests.unit_tests.config_override import apply_config_overrides
|
|
|
|
_NEEDS_METHOD_VIEW_CLEANUP = False
|
|
if not hasattr(builtins, "MethodView"):
|
|
builtins.__dict__["MethodView"] = FlaskMethodView
|
|
_NEEDS_METHOD_VIEW_CLEANUP = True
|
|
|
|
from constants import HIDDEN_VALUE
|
|
from controllers.console.extension import (
|
|
APIBasedExtensionAPI,
|
|
APIBasedExtensionDetailAPI,
|
|
CodeBasedExtensionAPI,
|
|
)
|
|
from enums import DeploymentEdition
|
|
|
|
if _NEEDS_METHOD_VIEW_CLEANUP:
|
|
del builtins.__dict__["MethodView"]
|
|
from models.account import Account, AccountStatus, Tenant
|
|
from models.api_based_extension import APIBasedExtension
|
|
|
|
|
|
def _make_extension(
|
|
*,
|
|
name: str = "Sample Extension",
|
|
api_endpoint: str = "https://example.com/api",
|
|
api_key: str = "super-secret-key",
|
|
) -> APIBasedExtension:
|
|
extension = APIBasedExtension(
|
|
tenant_id="tenant-123",
|
|
name=name,
|
|
api_endpoint=api_endpoint,
|
|
api_key=api_key,
|
|
)
|
|
extension.id = f"{uuid.uuid4()}"
|
|
extension.created_at = datetime.now(tz=UTC)
|
|
return extension
|
|
|
|
|
|
def _masked_api_key(api_key: str) -> str:
|
|
if len(api_key) <= 8:
|
|
return api_key[0] + "******" + api_key[-1]
|
|
return api_key[:3] + "******" + api_key[-3:]
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _mock_console_guards(monkeypatch: pytest.MonkeyPatch) -> Account:
|
|
"""Bypass console decorators so handlers can run in isolation."""
|
|
|
|
from controllers.console import wraps as wraps_module
|
|
|
|
tenant = Tenant(name="Test Workspace")
|
|
tenant.id = "tenant-123"
|
|
account = Account(name="Test User", email="user@example.com", status=AccountStatus.ACTIVE)
|
|
account.id = "account-123"
|
|
account._current_tenant = tenant
|
|
|
|
apply_config_overrides(
|
|
monkeypatch,
|
|
DEPLOYMENT_EDITION=DeploymentEdition.CLOUD,
|
|
INIT_PASSWORD="",
|
|
LOGIN_DISABLED=True,
|
|
)
|
|
monkeypatch.setattr(wraps_module, "current_account_with_tenant", lambda: (account, "tenant-123"))
|
|
|
|
# The login_required decorator consults the shared LocalProxy in libs.login.
|
|
monkeypatch.setattr("libs.login.current_user", account)
|
|
monkeypatch.setattr("libs.login.check_csrf_token", lambda *_, **__: None)
|
|
|
|
return account
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _restx_mask_defaults(app: Flask):
|
|
app.config.setdefault("RESTX_MASK_HEADER", "X-Fields")
|
|
app.config.setdefault("RESTX_MASK_SWAGGER", False)
|
|
|
|
|
|
def test_code_based_extension_get_returns_service_data(app: Flask, monkeypatch: pytest.MonkeyPatch):
|
|
service_result = [{"entrypoint": "main:agent"}]
|
|
service_mock = MagicMock(return_value=service_result)
|
|
monkeypatch.setattr(
|
|
"controllers.console.extension.CodeBasedExtensionService.get_code_based_extension",
|
|
service_mock,
|
|
)
|
|
|
|
with app.test_request_context(
|
|
"/console/api/code-based-extension",
|
|
method="GET",
|
|
query_string={"module": "workflow.tools"},
|
|
):
|
|
response = CodeBasedExtensionAPI().get()
|
|
|
|
assert response == {"module": "workflow.tools", "data": service_result}
|
|
service_mock.assert_called_once_with("workflow.tools")
|
|
|
|
|
|
def test_api_based_extension_get_returns_tenant_extensions(app: Flask, monkeypatch: pytest.MonkeyPatch):
|
|
extension = _make_extension(name="Weather API", api_key="abcdefghi123")
|
|
service_mock = MagicMock(return_value=[extension])
|
|
monkeypatch.setattr(
|
|
"controllers.console.extension.APIBasedExtensionService.get_all_by_tenant_id",
|
|
service_mock,
|
|
)
|
|
|
|
with app.test_request_context("/console/api/api-based-extension", method="GET"):
|
|
response = APIBasedExtensionAPI().get()
|
|
|
|
assert response[0]["id"] == extension.id
|
|
assert response[0]["name"] == "Weather API"
|
|
assert response[0]["api_endpoint"] == extension.api_endpoint
|
|
assert response[0]["api_key"].startswith(extension.api_key[:3])
|
|
service_mock.assert_called_once_with("tenant-123", session=ANY)
|
|
|
|
|
|
def test_api_based_extension_post_creates_extension(app: Flask, monkeypatch: pytest.MonkeyPatch):
|
|
saved_extension = _make_extension(name="Docs API", api_key="encrypted-token-from-save")
|
|
save_mock = MagicMock(return_value=saved_extension)
|
|
monkeypatch.setattr("controllers.console.extension.APIBasedExtensionService.save", save_mock)
|
|
|
|
payload = {
|
|
"name": "Docs API",
|
|
"api_endpoint": "https://docs.example.com/hook",
|
|
"api_key": "plain-secret",
|
|
}
|
|
|
|
with app.test_request_context("/console/api/api-based-extension", method="POST", json=payload):
|
|
response, status = APIBasedExtensionAPI().post()
|
|
|
|
args, _ = save_mock.call_args
|
|
created_extension: APIBasedExtension = args[0]
|
|
assert created_extension.tenant_id == "tenant-123"
|
|
assert created_extension.name == payload["name"]
|
|
assert created_extension.api_endpoint == payload["api_endpoint"]
|
|
assert created_extension.api_key == payload["api_key"]
|
|
assert status == 201
|
|
assert response["name"] == saved_extension.name
|
|
assert response["api_key"] == _masked_api_key(payload["api_key"])
|
|
save_mock.assert_called_once()
|
|
|
|
|
|
def test_api_based_extension_detail_get_fetches_extension(app: Flask, monkeypatch: pytest.MonkeyPatch):
|
|
extension = _make_extension(name="Docs API", api_key="abcdefg12345")
|
|
service_mock = MagicMock(return_value=extension)
|
|
monkeypatch.setattr(
|
|
"controllers.console.extension.APIBasedExtensionService.get_with_tenant_id",
|
|
service_mock,
|
|
)
|
|
|
|
extension_id = uuid.uuid4()
|
|
with app.test_request_context(f"/console/api/api-based-extension/{extension_id}", method="GET"):
|
|
response = APIBasedExtensionDetailAPI().get(extension_id)
|
|
|
|
assert response["id"] == extension.id
|
|
assert response["name"] == extension.name
|
|
service_mock.assert_called_once_with("tenant-123", str(extension_id), session=ANY)
|
|
|
|
|
|
def test_api_based_extension_detail_post_keeps_hidden_api_key(app: Flask, monkeypatch: pytest.MonkeyPatch):
|
|
existing_extension = _make_extension(name="Docs API", api_key="keep-me")
|
|
get_mock = MagicMock(return_value=existing_extension)
|
|
save_mock = MagicMock(return_value=existing_extension)
|
|
monkeypatch.setattr(
|
|
"controllers.console.extension.APIBasedExtensionService.get_with_tenant_id",
|
|
get_mock,
|
|
)
|
|
monkeypatch.setattr("controllers.console.extension.APIBasedExtensionService.save", save_mock)
|
|
|
|
payload = {
|
|
"name": "Docs API Updated",
|
|
"api_endpoint": "https://docs.example.com/v2",
|
|
"api_key": HIDDEN_VALUE,
|
|
}
|
|
|
|
extension_id = uuid.uuid4()
|
|
with app.test_request_context(
|
|
f"/console/api/api-based-extension/{extension_id}",
|
|
method="POST",
|
|
json=payload,
|
|
):
|
|
response = APIBasedExtensionDetailAPI().post(extension_id)
|
|
|
|
assert existing_extension.name == payload["name"]
|
|
assert existing_extension.api_endpoint == payload["api_endpoint"]
|
|
assert existing_extension.api_key == "keep-me"
|
|
save_mock.assert_called_once_with(existing_extension, session=ANY)
|
|
assert response["name"] == payload["name"]
|
|
assert response["api_key"] == _masked_api_key("keep-me")
|
|
|
|
|
|
def test_api_based_extension_detail_post_updates_api_key_when_provided(app: Flask, monkeypatch: pytest.MonkeyPatch):
|
|
existing_extension = _make_extension(name="Docs API", api_key="old-secret")
|
|
get_mock = MagicMock(return_value=existing_extension)
|
|
save_mock = MagicMock(return_value=existing_extension)
|
|
monkeypatch.setattr(
|
|
"controllers.console.extension.APIBasedExtensionService.get_with_tenant_id",
|
|
get_mock,
|
|
)
|
|
monkeypatch.setattr("controllers.console.extension.APIBasedExtensionService.save", save_mock)
|
|
|
|
payload = {
|
|
"name": "Docs API Updated",
|
|
"api_endpoint": "https://docs.example.com/v2",
|
|
"api_key": "new-secret",
|
|
}
|
|
|
|
extension_id = uuid.uuid4()
|
|
with app.test_request_context(
|
|
f"/console/api/api-based-extension/{extension_id}",
|
|
method="POST",
|
|
json=payload,
|
|
):
|
|
response = APIBasedExtensionDetailAPI().post(extension_id)
|
|
|
|
assert existing_extension.api_key == "new-secret"
|
|
save_mock.assert_called_once_with(existing_extension, session=ANY)
|
|
assert response["name"] == payload["name"]
|
|
assert response["api_key"] == _masked_api_key(payload["api_key"])
|
|
|
|
|
|
def test_api_based_extension_detail_delete_removes_extension(app: Flask, monkeypatch: pytest.MonkeyPatch):
|
|
existing_extension = _make_extension()
|
|
get_mock = MagicMock(return_value=existing_extension)
|
|
delete_mock = MagicMock()
|
|
monkeypatch.setattr(
|
|
"controllers.console.extension.APIBasedExtensionService.get_with_tenant_id",
|
|
get_mock,
|
|
)
|
|
monkeypatch.setattr("controllers.console.extension.APIBasedExtensionService.delete", delete_mock)
|
|
|
|
extension_id = uuid.uuid4()
|
|
with app.test_request_context(
|
|
f"/console/api/api-based-extension/{extension_id}",
|
|
method="DELETE",
|
|
):
|
|
response, status = APIBasedExtensionDetailAPI().delete(extension_id)
|
|
|
|
delete_mock.assert_called_once_with(existing_extension, session=ANY)
|
|
assert status == 204
|
|
assert response == ""
|