diff --git a/api/tests/unit_tests/controllers/inner_api/knowledge_fs/test_texts.py b/api/tests/unit_tests/controllers/inner_api/knowledge_fs/test_texts.py index 1c1bc180fa8..6136c153a87 100644 --- a/api/tests/unit_tests/controllers/inner_api/knowledge_fs/test_texts.py +++ b/api/tests/unit_tests/controllers/inner_api/knowledge_fs/test_texts.py @@ -1,7 +1,9 @@ import inspect +from collections.abc import Callable from unittest.mock import patch import pytest +from flask import Flask from werkzeug.exceptions import BadRequest, NotFound, ServiceUnavailable from controllers.inner_api.knowledge_fs.texts import KnowledgeFSTextApi @@ -17,13 +19,13 @@ PAYLOAD = { } -def test_text_controller_rejects_unauthenticated_requests(app): +def test_text_controller_rejects_unauthenticated_requests(app: Flask) -> None: with app.test_request_context(json=PAYLOAD): with pytest.raises(NotFound): KnowledgeFSTextApi().post() -def test_text_controller_validates_before_opening_backend(app): +def test_text_controller_validates_before_opening_backend(app: Flask) -> None: handler = KnowledgeFSTextApi() with ( app.test_request_context(json={**PAYLOAD, "endpoint": "http://untrusted"}), @@ -34,7 +36,7 @@ def test_text_controller_validates_before_opening_backend(app): client.assert_not_called() -def test_text_controller_hides_provider_credentials(app): +def test_text_controller_hides_provider_credentials(app: Flask) -> None: handler = KnowledgeFSTextApi() with ( app.test_request_context(json=PAYLOAD), @@ -49,9 +51,9 @@ def test_text_controller_hides_provider_credentials(app): assert error.value.__suppress_context__ -def test_text_controller_serializes_successful_response(app): +def test_text_controller_serializes_successful_response(app: Flask) -> None: handler = KnowledgeFSTextApi() - result = {"points": [], "matches": [{"id": PAYLOAD["ids"][0], "score": 0.75}]} + result: dict[str, list[dict[str, object]]] = {"points": [], "matches": [{"id": PAYLOAD["ids"][0], "score": 0.75}]} with ( app.test_request_context(json=PAYLOAD), patch("controllers.inner_api.knowledge_fs.texts.configured_text_client"), @@ -60,7 +62,7 @@ def test_text_controller_serializes_successful_response(app): assert inspect.unwrap(handler.post)(handler) == result -def test_text_controller_accepts_only_configured_inner_key(app, config_overrides): +def test_text_controller_accepts_only_configured_inner_key(app: Flask, config_overrides: Callable[..., None]) -> None: config_overrides(PLUGIN_DAEMON_KEY="enabled", INNER_API_KEY_FOR_PLUGIN="trusted-key") with ( patch("controllers.inner_api.knowledge_fs.texts.configured_text_client") as client, @@ -79,7 +81,7 @@ def test_text_controller_accepts_only_configured_inner_key(app, config_overrides client.assert_called_once_with(PAYLOAD["scope"]["tenant_id"], allow_create=False) -def test_text_controller_reports_pending_index_without_leaking_credentials(app): +def test_text_controller_reports_pending_index_without_leaking_credentials(app: Flask) -> None: from services.knowledge_fs.text_store import TextIndexPendingError handler = KnowledgeFSTextApi() diff --git a/api/tests/unit_tests/controllers/inner_api/knowledge_fs/test_vectors.py b/api/tests/unit_tests/controllers/inner_api/knowledge_fs/test_vectors.py index 291961966a9..daa7530cc0a 100644 --- a/api/tests/unit_tests/controllers/inner_api/knowledge_fs/test_vectors.py +++ b/api/tests/unit_tests/controllers/inner_api/knowledge_fs/test_vectors.py @@ -1,8 +1,10 @@ import inspect -from contextlib import nullcontext +from collections.abc import Callable +from contextlib import AbstractContextManager, nullcontext from unittest.mock import patch import pytest +from flask import Flask from werkzeug.exceptions import BadRequest, NotFound, ServiceUnavailable from controllers.inner_api.knowledge_fs.vectors import KnowledgeFSVectorApi, KnowledgeFSVectorQuotaExceededError @@ -21,13 +23,13 @@ PAYLOAD = { } -def test_vector_controller_rejects_unauthenticated_requests(app): +def test_vector_controller_rejects_unauthenticated_requests(app: Flask) -> None: with app.test_request_context(json=PAYLOAD): with pytest.raises(NotFound): KnowledgeFSVectorApi().post() -def test_vector_controller_validates_before_opening_backend(app): +def test_vector_controller_validates_before_opening_backend(app: Flask) -> None: handler = KnowledgeFSVectorApi() with ( app.test_request_context(json={**PAYLOAD, "endpoint": "http://untrusted"}), @@ -38,7 +40,7 @@ def test_vector_controller_validates_before_opening_backend(app): client.assert_not_called() -def test_vector_controller_hides_provider_credentials(app): +def test_vector_controller_hides_provider_credentials(app: Flask) -> None: handler = KnowledgeFSVectorApi() with ( app.test_request_context(json=PAYLOAD), @@ -53,9 +55,9 @@ def test_vector_controller_hides_provider_credentials(app): assert error.value.__suppress_context__ -def test_vector_controller_serializes_successful_response(app): +def test_vector_controller_serializes_successful_response(app: Flask) -> None: handler = KnowledgeFSVectorApi() - result = {"points": [], "matches": [{"id": PAYLOAD["ids"][0], "score": 0.75}]} + result: dict[str, list[dict[str, object]]] = {"points": [], "matches": [{"id": PAYLOAD["ids"][0], "score": 0.75}]} with ( app.test_request_context(json=PAYLOAD), patch("controllers.inner_api.knowledge_fs.vectors.configured_vector_client"), @@ -64,7 +66,7 @@ def test_vector_controller_serializes_successful_response(app): assert inspect.unwrap(handler.post)(handler) == result -def test_vector_controller_accepts_only_configured_inner_key(app, config_overrides): +def test_vector_controller_accepts_only_configured_inner_key(app: Flask, config_overrides: Callable[..., None]) -> None: config_overrides(PLUGIN_DAEMON_KEY="enabled", INNER_API_KEY_FOR_PLUGIN="trusted-key") with ( patch("controllers.inner_api.knowledge_fs.vectors.configured_vector_client") as client, @@ -83,7 +85,7 @@ def test_vector_controller_accepts_only_configured_inner_key(app, config_overrid client.assert_called_once_with(PAYLOAD["scope"]["tenant_id"], allow_create=False) -def test_vector_controller_reports_quota_failure(app): +def test_vector_controller_reports_quota_failure(app: Flask) -> None: handler = KnowledgeFSVectorApi() with ( app.test_request_context(json=PAYLOAD), @@ -98,9 +100,9 @@ def test_vector_controller_reports_quota_failure(app): @pytest.mark.parametrize("operation", ["get", "search", "delete", "upsert"]) -def test_only_admitted_writes_can_provision(app, operation): +def test_only_admitted_writes_can_provision(app: Flask, operation: str) -> None: handler = KnowledgeFSVectorApi() - payload = {**PAYLOAD, "operation": operation} + payload: dict[str, object] = {**PAYLOAD, "operation": operation} if operation == "search": payload["query_vector"] = [1, 0] if operation == "upsert": @@ -118,7 +120,7 @@ def test_only_admitted_writes_can_provision(app, operation): ) as execute, ): - def open_client(*_args, **_kwargs): + def open_client(*_args: object, **_kwargs: object) -> AbstractContextManager[object]: admit.assert_called_once() return nullcontext(object()) @@ -128,7 +130,7 @@ def test_only_admitted_writes_can_provision(app, operation): assert execute.call_args.kwargs["check_admission"] is False -def test_quota_rejection_cannot_allocate_a_cluster(app): +def test_quota_rejection_cannot_allocate_a_cluster(app: Flask) -> None: handler = KnowledgeFSVectorApi() with ( app.test_request_context(json=PAYLOAD), @@ -150,8 +152,9 @@ def test_quota_rejection_cannot_allocate_a_cluster(app): (VectorSpaceAdmissionUnavailableError("vector usage is unavailable"), 503, "vector_space_quota_unavailable"), ], ) -def test_http_quota_failure_cannot_allocate_a_cluster(config_overrides, admission_error, status_code, error_code): - from flask import Flask +def test_http_quota_failure_cannot_allocate_a_cluster( + config_overrides: Callable[..., None], admission_error: Exception, status_code: int, error_code: str +) -> None: from flask_restx import Api config_overrides(PLUGIN_DAEMON_KEY="enabled", INNER_API_KEY_FOR_PLUGIN="trusted-key") @@ -177,7 +180,7 @@ def test_http_quota_failure_cannot_allocate_a_cluster(config_overrides, admissio execute.assert_not_called() -def test_pending_cluster_returns_explicit_retry_after_without_secrets(app): +def test_pending_cluster_returns_explicit_retry_after_without_secrets(app: Flask) -> None: from services.tidb_binding_service import TidbBindingPendingError handler = KnowledgeFSVectorApi() @@ -194,8 +197,7 @@ def test_pending_cluster_returns_explicit_retry_after_without_secrets(app): assert "private" not in str(error.value) -def test_http_pending_response_preserves_retry_after_header(config_overrides): - from flask import Flask +def test_http_pending_response_preserves_retry_after_header(config_overrides: Callable[..., None]) -> None: from flask_restx import Api from services.tidb_binding_service import TidbBindingPendingError diff --git a/api/tests/unit_tests/controllers/inner_api/test_agent_knowledge_investigations.py b/api/tests/unit_tests/controllers/inner_api/test_agent_knowledge_investigations.py index 7abaf564df6..ea5711a4eea 100644 --- a/api/tests/unit_tests/controllers/inner_api/test_agent_knowledge_investigations.py +++ b/api/tests/unit_tests/controllers/inner_api/test_agent_knowledge_investigations.py @@ -1,11 +1,14 @@ """Queue admission and authorization boundaries for runtime quality reports.""" -from unittest.mock import patch +from collections.abc import Iterator +from typing import TypedDict +from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest from dify_agent.protocol.knowledge_fs import KnowledgeFsError from flask import Flask +from flask.testing import FlaskClient from flask_restx.swagger import Swagger from controllers.inner_api import api as inner_api @@ -16,7 +19,16 @@ PATH = "/inner/api/agent/knowledge/investigations" HEADERS = {"X-Inner-Api-Key": "inner-key"} -def _report() -> dict: +class _Report(TypedDict): + investigation_id: str + execution_context: dict[str, str] + bindings: list[dict[str, str]] + query: str + status: str + attempts: list[dict[str, object]] + + +def _report() -> _Report: spaces = [str(uuid4()), str(uuid4())] return { "investigation_id": str(uuid4()), @@ -42,7 +54,7 @@ def _report() -> dict: @pytest.fixture -def boundary(): +def boundary() -> Iterator[tuple[FlaskClient, MagicMock, MagicMock]]: app = Flask(__name__) app.config["TESTING"] = True app.register_blueprint(inner_api_bp) @@ -57,7 +69,7 @@ def boundary(): yield app.test_client(), gateway.return_value, enqueue -def test_queues_once_per_space_with_original_identity(boundary) -> None: +def test_queues_once_per_space_with_original_identity(boundary: tuple[FlaskClient, MagicMock, MagicMock]) -> None: client, gateway, enqueue = boundary payload = _report() response = client.post(PATH, json=payload, headers=HEADERS) @@ -75,27 +87,27 @@ def test_queues_once_per_space_with_original_identity(boundary) -> None: gateway.capture_investigation.assert_not_called() -def test_revoked_scope_cannot_enqueue_a_report(boundary) -> None: +def test_revoked_scope_cannot_enqueue_a_report(boundary: tuple[FlaskClient, MagicMock, MagicMock]) -> None: client, gateway, enqueue = boundary gateway.authorize_investigation.side_effect = KnowledgeFsError("KNOWLEDGE_SCOPE_MISMATCH", "revoked", 403) assert client.post(PATH, json=_report(), headers=HEADERS).status_code == 403 enqueue.assert_not_called() -def test_broker_failure_is_not_acknowledged_as_accepted(boundary) -> None: +def test_broker_failure_is_not_acknowledged_as_accepted(boundary: tuple[FlaskClient, MagicMock, MagicMock]) -> None: client, _, enqueue = boundary enqueue.side_effect = OSError("broker offline") assert client.post(PATH, json=_report(), headers=HEADERS).status_code == 500 -def test_missing_inner_api_key_cannot_enqueue_a_report(boundary) -> None: +def test_missing_inner_api_key_cannot_enqueue_a_report(boundary: tuple[FlaskClient, MagicMock, MagicMock]) -> None: client, gateway, enqueue = boundary assert client.post(PATH, json=_report()).status_code == 404 gateway.authorize_investigation.assert_not_called() enqueue.assert_not_called() -def test_swagger_documents_report_and_queue_admission(boundary) -> None: +def test_swagger_documents_report_and_queue_admission(boundary: tuple[FlaskClient, MagicMock, MagicMock]) -> None: client, _, _ = boundary with client.application.test_request_context(): schema = Swagger(inner_api).as_dict() diff --git a/api/tests/unit_tests/controllers/inner_api/test_auth_wraps.py b/api/tests/unit_tests/controllers/inner_api/test_auth_wraps.py index a5b6b0f7fae..07c106312a6 100644 --- a/api/tests/unit_tests/controllers/inner_api/test_auth_wraps.py +++ b/api/tests/unit_tests/controllers/inner_api/test_auth_wraps.py @@ -311,18 +311,18 @@ class TestPluginInnerApiOnly: class TestKnowledgeFSInnerApiOnly: """KnowledgeFS uses the same trusted transport key without exposing plugin semantics.""" - def test_should_allow_valid_shared_inner_key(self, app: Flask, config_overrides): + def test_should_allow_valid_shared_inner_key(self, app: Flask, config_overrides: Callable[..., None]) -> None: @knowledge_fs_inner_api_only - def protected_view(): + def protected_view() -> str: return "success" config_overrides(PLUGIN_DAEMON_KEY="plugin_key", INNER_API_KEY_FOR_PLUGIN="valid_plugin_key") with app.test_request_context(headers={"X-Inner-Api-Key": "valid_plugin_key"}): assert protected_view() == "success" - def test_should_hide_endpoint_for_invalid_key(self, app: Flask, config_overrides): + def test_should_hide_endpoint_for_invalid_key(self, app: Flask, config_overrides: Callable[..., None]) -> None: @knowledge_fs_inner_api_only - def protected_view(): + def protected_view() -> str: return "success" config_overrides(PLUGIN_DAEMON_KEY="plugin_key", INNER_API_KEY_FOR_PLUGIN="valid_plugin_key")