chore(api): Fix several typing errors (#37248)

api/tests/unit_tests/core/tools/test_builtin_tool_provider.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Generator
-from typing import Any
+from typing import Any, override
 from unittest.mock import patch
 
 import pytest
@@ -16,6 +16,7 @@ from core.tools.errors import ToolProviderNotFoundError
 
 
 class _FakeBuiltinTool(BuiltinTool):
+    @override
     def _invoke(
         self,
         user_id: str,
@@ -30,6 +31,7 @@ class _FakeBuiltinTool(BuiltinTool):
 class _ConcreteBuiltinProvider(BuiltinToolProviderController):
     last_validation: tuple[str, dict[str, Any]] | None = None
 
+    @override
     def _validate_credentials(self, user_id: str, credentials: dict[str, Any]):
         self.last_validation = (user_id, credentials)
 

api/tests/unit_tests/core/tools/test_tool_label_manager.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from types import SimpleNamespace
-from typing import Any
+from typing import Any, override
 from unittest.mock import MagicMock, PropertyMock, patch
 
 import pytest
@@ -14,6 +14,7 @@ from core.tools.workflow_as_tool.provider import WorkflowToolProviderController
 
 # Create a mock class for testing abstract/base classes
 class _ConcreteBuiltinToolProviderController(BuiltinToolProviderController):
+    @override
     def _validate_credentials(self, user_id: str, credentials: dict[str, Any]):
         return None
 

api/tests/unit_tests/core/tools/test_tool_provider_controller.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Generator
-from typing import Any
+from typing import Any, override
 
 import pytest
 
@@ -22,9 +22,11 @@ from core.tools.errors import ToolProviderCredentialValidationError
 
 
 class _DummyTool(Tool):
+    @override
     def tool_provider_type(self) -> ToolProviderType:
         return ToolProviderType.BUILT_IN
 
+    @override
     def _invoke(
         self,
         user_id: str,
@@ -36,7 +38,8 @@ class _DummyTool(Tool):
         yield self.create_text_message("ok")
 
 
-class _DummyController(ToolProviderController):
+class _DummyController(ToolProviderController[ToolProviderEntity, Tool]):
+    @override
     def get_tool(self, tool_name: str) -> Tool:
         entity = ToolEntity(
             identity=ToolIdentity(

api/tests/unit_tests/core/tools/workflow_as_tool/test_provider.py

@@ -1,19 +1,30 @@
 from __future__ import annotations
 
+import json
 from types import SimpleNamespace
+from typing import Any, cast
 from unittest.mock import MagicMock, Mock, patch
 
 import pytest
 
+from core.tools.__base.tool_runtime import ToolRuntime
 from core.tools.entities.common_entities import I18nObject
 from core.tools.entities.tool_entities import (
+    ToolDescription,
+    ToolEntity,
+    ToolIdentity,
     ToolParameter,
     ToolProviderEntity,
     ToolProviderIdentity,
     ToolProviderType,
 )
 from core.tools.workflow_as_tool.provider import WorkflowToolProviderController
+from core.tools.workflow_as_tool.tool import WorkflowTool
 from graphon.variables.input_entities import VariableEntity, VariableEntityType
+from models.account import Account
+from models.model import App
+from models.tools import WorkflowToolProvider
+from models.workflow import Workflow, WorkflowType
 
 
 def _controller() -> WorkflowToolProviderController:
@@ -30,6 +41,64 @@ def _controller() -> WorkflowToolProviderController:
     return WorkflowToolProviderController(entity=entity, provider_id="provider-1")
 
 
+def _app() -> App:
+    return App(id="app-1")
+
+
+def _account() -> Account:
+    return Account(name="Alice", email="alice@example.com")
+
+
+def _workflow() -> Workflow:
+    return Workflow.new(
+        tenant_id="tenant-1",
+        app_id="app-1",
+        type=WorkflowType.WORKFLOW.value,
+        version="1",
+        graph=json.dumps({"nodes": []}),
+        features="{}",
+        created_by="user-1",
+        environment_variables=[],
+        conversation_variables=[],
+        rag_pipeline_variables=[],
+    )
+
+
+def _db_provider(*, parameter_configuration: str = "[]") -> WorkflowToolProvider:
+    return WorkflowToolProvider(
+        name="workflow_tool",
+        label="WF Provider",
+        icon="icon.svg",
+        app_id="app-1",
+        version="1",
+        user_id="user-1",
+        tenant_id="tenant-1",
+        description="desc",
+        parameter_configuration=parameter_configuration,
+    )
+
+
+def _workflow_tool(name: str = "workflow_tool") -> WorkflowTool:
+    return WorkflowTool(
+        workflow_as_tool_id="provider-1",
+        entity=ToolEntity(
+            identity=ToolIdentity(
+                author="author",
+                name=name,
+                label=I18nObject(en_US=name),
+                provider="provider-1",
+            ),
+            description=ToolDescription(human=I18nObject(en_US="desc"), llm="desc"),
+            parameters=[],
+        ),
+        runtime=ToolRuntime(tenant_id="tenant-1"),
+        workflow_app_id="app-1",
+        workflow_entities={"app": _app(), "workflow": _workflow()},
+        version="1",
+        workflow_call_depth=0,
+    )
+
+
 def _mock_session_with_begin() -> Mock:
     session = Mock()
     begin_cm = Mock()
@@ -42,25 +111,18 @@ def _mock_session_with_begin() -> Mock:
 def test_get_db_provider_tool_builds_entity():
     controller = _controller()
     session = Mock()
-    workflow = SimpleNamespace(graph_dict={"nodes": []}, features_dict={})
+    workflow = _workflow()
     session.scalar.return_value = workflow
-    app = SimpleNamespace(id="app-1")
-    db_provider = SimpleNamespace(
-        id="provider-1",
-        app_id="app-1",
-        version="1",
-        label="WF Provider",
-        description="desc",
-        icon="icon.svg",
-        name="workflow_tool",
-        tenant_id="tenant-1",
-        user_id="user-1",
-        parameter_configurations=[
-            SimpleNamespace(name="country", description="Country", form=ToolParameter.ToolParameterForm.FORM),
-            SimpleNamespace(name="files", description="files", form=ToolParameter.ToolParameterForm.FORM),
-        ],
+    app = _app()
+    db_provider = _db_provider(
+        parameter_configuration=json.dumps(
+            [
+                {"name": "country", "description": "Country", "form": ToolParameter.ToolParameterForm.FORM.value},
+                {"name": "files", "description": "files", "form": ToolParameter.ToolParameterForm.FORM.value},
+            ]
+        )
     )
-    user = SimpleNamespace(name="Alice")
+    user = _account()
     variables = [
         VariableEntity(
             variable="country",
@@ -94,8 +156,9 @@ def test_get_db_provider_tool_builds_entity():
 
     assert tool.entity.identity.name == "workflow_tool"
     # "json" output is reserved for ToolInvokeMessage.VariableMessage and filtered out.
-    assert tool.entity.output_schema["properties"] == {"answer": {"type": "string", "description": ""}}
-    assert "json" not in tool.entity.output_schema["properties"]
+    properties = cast(dict[str, Any], tool.entity.output_schema["properties"])
+    assert properties == {"answer": {"type": "string", "description": ""}}
+    assert "json" not in properties
     assert tool.entity.parameters[0].type == ToolParameter.ToolParameterType.SELECT
     assert tool.entity.parameters[1].type == ToolParameter.ToolParameterType.SYSTEM_FILES
     assert controller.provider_type == ToolProviderType.WORKFLOW
@@ -103,7 +166,7 @@ def test_get_db_provider_tool_builds_entity():
 
 def test_get_tool_returns_hit_or_none():
     controller = _controller()
-    tool = SimpleNamespace(entity=SimpleNamespace(identity=SimpleNamespace(name="workflow_tool")))
+    tool = _workflow_tool()
     controller.tools = [tool]
 
     assert controller.get_tool("workflow_tool") is tool
@@ -112,29 +175,16 @@ def test_get_tool_returns_hit_or_none():
 
 def test_get_tools_returns_cached():
     controller = _controller()
-    cached_tools = [SimpleNamespace(entity=SimpleNamespace(identity=SimpleNamespace(name="wf-cached")))]
-    controller.tools = cached_tools  # type: ignore[assignment]
+    cached_tools = [_workflow_tool("wf-cached")]
+    controller.tools = cached_tools
 
     assert controller.get_tools("tenant-1") == cached_tools
 
 
 def test_from_db_builds_controller():
-    controller = _controller()
-
-    app = SimpleNamespace(id="app-1")
-    user = SimpleNamespace(name="Alice")
-    db_provider = SimpleNamespace(
-        id="provider-1",
-        app_id="app-1",
-        version="1",
-        user_id="user-1",
-        label="WF Provider",
-        description="desc",
-        icon="icon.svg",
-        name="workflow_tool",
-        tenant_id="tenant-1",
-        parameter_configurations=[],
-    )
+    app = _app()
+    user = _account()
+    db_provider = _db_provider()
     session = _mock_session_with_begin()
     session.scalar.return_value = db_provider
     session.get.side_effect = [app, user]
@@ -148,7 +198,7 @@ def test_from_db_builds_controller():
         with patch.object(
             WorkflowToolProviderController,
             "_get_db_provider_tool",
-            return_value=SimpleNamespace(entity=SimpleNamespace(identity=SimpleNamespace(name="wf"))),
+            return_value=_workflow_tool("wf"),
         ):
             built = WorkflowToolProviderController.from_db(db_provider)
     assert isinstance(built, WorkflowToolProviderController)
@@ -157,7 +207,7 @@ def test_from_db_builds_controller():
 
 def test_get_tools_returns_empty_when_provider_missing():
     controller = _controller()
-    controller.tools = None  # type: ignore[assignment]
+    controller.tools = None
 
     with patch("core.tools.workflow_as_tool.provider.db") as mock_db:
         mock_db.engine = object()
@@ -171,19 +221,8 @@ def test_get_tools_returns_empty_when_provider_missing():
 
 def test_get_tools_raises_when_app_missing():
     controller = _controller()
-    controller.tools = None  # type: ignore[assignment]
-    db_provider = SimpleNamespace(
-        id="provider-1",
-        app_id="app-1",
-        version="1",
-        user_id="user-1",
-        label="WF Provider",
-        description="desc",
-        icon="icon.svg",
-        name="workflow_tool",
-        tenant_id="tenant-1",
-        parameter_configurations=[],
-    )
+    controller.tools = None
+    db_provider = _db_provider()
 
     with patch("core.tools.workflow_as_tool.provider.db") as mock_db:
         mock_db.engine = object()

api/tests/unit_tests/services/test_document_indexing_task_proxy.py

@@ -10,7 +10,7 @@ class DocumentIndexingTaskProxyTestDataFactory:
     """Factory class for creating test data and mock objects for DocumentIndexingTaskProxy tests."""
 
     @staticmethod
-    def create_mock_features(billing_enabled: bool = False, plan: CloudPlan = CloudPlan.SANDBOX) -> Mock:
+    def create_mock_features(billing_enabled: bool = False, plan: CloudPlan | str | None = CloudPlan.SANDBOX) -> Mock:
         """Create mock features with billing configuration."""
         features = Mock()
         features.billing = Mock()

api/tests/unit_tests/services/test_duplicate_document_indexing_task_proxy.py

@@ -12,7 +12,7 @@ class DuplicateDocumentIndexingTaskProxyTestDataFactory:
     """Factory class for creating test data and mock objects for DuplicateDocumentIndexingTaskProxy tests."""
 
     @staticmethod
-    def create_mock_features(billing_enabled: bool = False, plan: CloudPlan = CloudPlan.SANDBOX) -> Mock:
+    def create_mock_features(billing_enabled: bool = False, plan: CloudPlan | str | None = CloudPlan.SANDBOX) -> Mock:
         """Create mock features with billing configuration."""
         features = Mock()
         features.billing = Mock()