refactor(api): enable strict typing for events/extensions test files (#44875)

Co-authored-by: Felix Shi <204130665+shi-felix@users.noreply.github.com>

api/tests/unit_tests/events/test___init__.py

@@ -3,12 +3,14 @@ import pytest
 from events import Events, EventsError, EventsException
 
 
-def test_events_package_exposes_opensearchpy_compatible_events_class():
+def test_events_package_exposes_opensearchpy_compatible_events_class() -> None:
     calls: list[str] = []
     events = Events()
 
-    events.request_start += lambda: calls.append("start")
-    events.request_end += lambda: calls.append("end")
+    # Augmented assignment writes the slot back onto the instance, so pyrefly cannot
+    # resolve it through `__getattr__`; plain reads below need no ignore.
+    events.request_start += lambda: calls.append("start")  # type: ignore[attr-defined]
+    events.request_end += lambda: calls.append("end")  # type: ignore[attr-defined]
 
     events.request_start()
     events.request_end()
@@ -16,7 +18,7 @@ def test_events_package_exposes_opensearchpy_compatible_events_class():
     assert calls == ["start", "end"]
 
 
-def test_events_package_supports_named_slots_iteration_removal_and_private_attrs():
+def test_events_package_supports_named_slots_iteration_removal_and_private_attrs() -> None:
     calls: list[str] = []
 
     def handler() -> None:
@@ -24,19 +26,19 @@ def test_events_package_supports_named_slots_iteration_removal_and_private_attrs
 
     events = Events("request_start")
 
-    events.request_start += handler
-    events.request_start += handler
+    events.request_start += handler  # type: ignore[attr-defined]
+    events.request_start += handler  # type: ignore[attr-defined]
 
     assert len(events.request_start) == 2
     assert list(events.request_start) == [handler, handler]
 
-    events.request_start -= handler
+    events.request_start -= handler  # type: ignore[attr-defined]
 
     assert len(events.request_start) == 0
     events.request_start()
     assert calls == []
 
     with pytest.raises(AttributeError):
-        _ = events._private  # type: ignore[attr-defined]
+        _ = events._private
 
     assert EventsException is EventsError

api/tests/unit_tests/events/test_app_event.py

@@ -5,11 +5,12 @@ from uuid import uuid4
 
 import pytest
 from sqlalchemy import event, select
-from sqlalchemy.orm import Session
+from sqlalchemy.orm import Session, sessionmaker
 
 from events.app_event import app_was_created, app_was_deleted, app_was_updated
 from events.event_handlers.create_installed_app_when_app_created import handle as create_installed_app
 from events.event_handlers.create_site_record_when_app_created import handle as create_site
+from extensions.application_services.app import AppServices
 from machinery.context import RequestContext
 from models.account import Account
 from models.agent import Agent, AgentWorkspace
@@ -67,7 +68,7 @@ def _make_collector[T](target: list[T]) -> Callable[..., None]:
 @pytest.mark.parametrize("sqlite_session", [(App, Account, Agent, AgentWorkspace)], indirect=True)
 @pytest.mark.usefixtures("_mock_deps")
 class TestAppWasDeletedSignal:
-    def test_sends_signal(self, app_services, app_model: App, sqlite_session: Session) -> None:
+    def test_sends_signal(self, app_services: AppServices, app_model: App, sqlite_session: Session) -> None:
         received: list[AppEvent] = []
         handler = _make_collector(received)
         app_was_deleted.connect(handler)
@@ -82,7 +83,11 @@ class TestAppWasDeletedSignal:
         assert sqlite_session.get(App, app_id) is None
 
     def test_signal_fires_after_db_commit(
-        self, sqlite_session_factory, app_services, app_model: App, sqlite_session: Session
+        self,
+        sqlite_session_factory: sessionmaker[Session],
+        app_services: AppServices,
+        app_model: App,
+        sqlite_session: Session,
     ) -> None:
         call_order: list[str] = []
 
@@ -115,7 +120,9 @@ class TestAppWasDeletedSignal:
 
 @pytest.mark.parametrize("sqlite_session", [(App, Account)], indirect=True)
 class TestAppWasUpdatedSignal:
-    def test_update_app(self, app_services, app_model: App, account: Account, sqlite_session: Session) -> None:
+    def test_update_app(
+        self, app_services: AppServices, app_model: App, account: Account, sqlite_session: Session
+    ) -> None:
         received: list[App] = []
         handler = _make_collector(received)
         app_was_updated.connect(handler)
@@ -149,7 +156,7 @@ class TestAppWasUpdatedSignal:
         assert persisted.updated_by == account.id
 
     def test_update_app_api_status_sends_when_changed(
-        self, app_services, app_model: App, account: Account, sqlite_session: Session
+        self, app_services: AppServices, app_model: App, account: Account, sqlite_session: Session
     ) -> None:
         received: list[App] = []
         handler = _make_collector(received)
@@ -167,7 +174,7 @@ class TestAppWasUpdatedSignal:
         assert sqlite_session.get(App, app_model.id).enable_api is True  # type: ignore[union-attr]
 
     def test_update_app_api_status_skips_when_unchanged(
-        self, app_services, app_model: App, sqlite_session: Session
+        self, app_services: AppServices, app_model: App, sqlite_session: Session
     ) -> None:
         app_model.enable_api = True
         sqlite_session.commit()
@@ -190,11 +197,11 @@ class TestAppModelConfigWasUpdatedSignal:
         from events.event_handlers.update_app_dataset_join_when_app_model_config_updated import handle
 
         with pytest.raises(TypeError, match="session"):
-            handle(app_model, app_model_config=None)
+            handle(app_model, app_model_config=None)  # pyrefly: ignore[missing-argument]
 
     @pytest.mark.parametrize("sqlite_session", [(App, Account, AppModelConfig, AppDatasetJoin)], indirect=True)
     def test_reuses_provided_session_without_committing(
-        self, sqlite_session_factory, app_model: App, sqlite_session: Session
+        self, sqlite_session_factory: sessionmaker[Session], app_model: App, sqlite_session: Session
     ) -> None:
         from events.event_handlers.update_app_dataset_join_when_app_model_config_updated import handle
 
@@ -280,7 +287,7 @@ class TestCreateInstalledAppWhenAppCreated:
         assert installations == [existing]
 
     def test_adds_missing_installation_without_committing(
-        self, sqlite_session_factory, app_model: App, sqlite_session: Session
+        self, sqlite_session_factory: sessionmaker[Session], app_model: App, sqlite_session: Session
     ) -> None:
         from events.event_handlers.create_installed_app_when_app_created import handle
 

api/tests/unit_tests/extensions/logstore/test_sql_escape.py

@@ -13,29 +13,29 @@ from extensions.logstore.sql_escape import escape_identifier, escape_logstore_qu
 class TestEscapeSQLString:
     """Test escape_sql_string function."""
 
-    def test_escape_empty_string(self):
+    def test_escape_empty_string(self) -> None:
         """Test escaping empty string."""
         assert escape_sql_string("") == ""
 
-    def test_escape_normal_string(self):
+    def test_escape_normal_string(self) -> None:
         """Test escaping string without special characters."""
         assert escape_sql_string("tenant_abc123") == "tenant_abc123"
         assert escape_sql_string("app-uuid-1234") == "app-uuid-1234"
 
-    def test_escape_single_quote(self):
+    def test_escape_single_quote(self) -> None:
         """Test escaping single quote."""
         # Single quote should be doubled
         assert escape_sql_string("tenant'id") == "tenant''id"
         assert escape_sql_string("O'Reilly") == "O''Reilly"
 
-    def test_escape_multiple_quotes(self):
+    def test_escape_multiple_quotes(self) -> None:
         """Test escaping multiple single quotes."""
         assert escape_sql_string("a'b'c") == "a''b''c"
         assert escape_sql_string("'''") == "''''''"
 
     # === SQL Injection Attack Scenarios ===
 
-    def test_prevent_boolean_injection(self):
+    def test_prevent_boolean_injection(self) -> None:
         """Test prevention of boolean injection attacks."""
         # Classic OR 1=1 attack
         malicious_input = "tenant' OR '1'='1"
@@ -47,7 +47,7 @@ class TestEscapeSQLString:
         assert sql == "WHERE tenant_id='tenant'' OR ''1''=''1'"
         # The entire input is now a string literal that won't match any tenant
 
-    def test_prevent_or_injection(self):
+    def test_prevent_or_injection(self) -> None:
         """Test prevention of OR-based injection."""
         malicious_input = "tenant_a' OR tenant_id='tenant_b"
         escaped = escape_sql_string(malicious_input)
@@ -59,7 +59,7 @@ class TestEscapeSQLString:
         # The SQL has: opening ', doubled internal quotes '', and closing '
         assert sql == "WHERE tenant_id='tenant_a'' OR tenant_id=''tenant_b'"
 
-    def test_prevent_union_injection(self):
+    def test_prevent_union_injection(self) -> None:
         """Test prevention of UNION-based injection."""
         malicious_input = "xxx' UNION SELECT password FROM users WHERE '1'='1"
         escaped = escape_sql_string(malicious_input)
@@ -69,7 +69,7 @@ class TestEscapeSQLString:
         assert "UNION" in escaped
         assert escaped.count("''") == 4  # All internal quotes are doubled
 
-    def test_prevent_comment_injection(self):
+    def test_prevent_comment_injection(self) -> None:
         """Test prevention of comment-based injection."""
         # SQL comment to bypass remaining conditions
         malicious_input = "tenant' --"
@@ -81,7 +81,7 @@ class TestEscapeSQLString:
         assert "--" in sql
         assert "AND deleted=false" in sql  # This part is NOT commented out
 
-    def test_prevent_semicolon_injection(self):
+    def test_prevent_semicolon_injection(self) -> None:
         """Test prevention of semicolon-based multi-statement injection."""
         malicious_input = "tenant'; DROP TABLE users; --"
         escaped = escape_sql_string(malicious_input)
@@ -90,7 +90,7 @@ class TestEscapeSQLString:
         # Semicolons and DROP are now part of the string
         assert "DROP TABLE" in escaped
 
-    def test_prevent_time_based_blind_injection(self):
+    def test_prevent_time_based_blind_injection(self) -> None:
         """Test prevention of time-based blind SQL injection."""
         malicious_input = "tenant' AND SLEEP(5) --"
         escaped = escape_sql_string(malicious_input)
@@ -99,7 +99,7 @@ class TestEscapeSQLString:
         # SLEEP becomes part of the string
         assert "SLEEP" in escaped
 
-    def test_prevent_wildcard_injection(self):
+    def test_prevent_wildcard_injection(self) -> None:
         """Test prevention of wildcard-based injection."""
         malicious_input = "tenant' OR tenant_id LIKE '%"
         escaped = escape_sql_string(malicious_input)
@@ -108,7 +108,7 @@ class TestEscapeSQLString:
         # The LIKE and wildcard are now part of the string
         assert "LIKE" in escaped
 
-    def test_prevent_null_byte_injection(self):
+    def test_prevent_null_byte_injection(self) -> None:
         """Test handling of null bytes."""
         # Null bytes can sometimes bypass filters
         malicious_input = "tenant\x00' OR '1'='1"
@@ -118,7 +118,7 @@ class TestEscapeSQLString:
 
     # === Real-world SAAS Scenarios ===
 
-    def test_cross_tenant_access_attempt(self):
+    def test_cross_tenant_access_attempt(self) -> None:
         """Test prevention of cross-tenant data access."""
         # Attacker tries to access another tenant's data
         attacker_input = "tenant_b' OR tenant_id='tenant_a"
@@ -129,7 +129,7 @@ class TestEscapeSQLString:
         # which doesn't exist - preventing access to either tenant's data
         assert "tenant_b'' OR tenant_id=''tenant_a" in sql
 
-    def test_cross_app_access_attempt(self):
+    def test_cross_app_access_attempt(self) -> None:
         """Test prevention of cross-application data access."""
         attacker_input = "app1' OR app_id='app2"
         escaped = escape_sql_string(attacker_input)
@@ -138,7 +138,7 @@ class TestEscapeSQLString:
         # Cannot access app2's data
         assert "app1'' OR app_id=''app2" in sql
 
-    def test_bypass_status_filter(self):
+    def test_bypass_status_filter(self) -> None:
         """Test prevention of bypassing status filters."""
         # Try to see all statuses instead of just 'running'
         attacker_input = "running' OR status LIKE '%"
@@ -150,18 +150,18 @@ class TestEscapeSQLString:
 
     # === Edge Cases ===
 
-    def test_escape_only_quotes(self):
+    def test_escape_only_quotes(self) -> None:
         """Test string with only quotes."""
         assert escape_sql_string("'") == "''"
         assert escape_sql_string("''") == "''''"
 
-    def test_escape_mixed_content(self):
+    def test_escape_mixed_content(self) -> None:
         """Test string with mixed quotes and other chars."""
         input_str = "It's a 'test' of O'Reilly's code"
         escaped = escape_sql_string(input_str)
         assert escaped == "It''s a ''test'' of O''Reilly''s code"
 
-    def test_escape_unicode_with_quotes(self):
+    def test_escape_unicode_with_quotes(self) -> None:
         """Test Unicode strings with quotes."""
         input_str = "租户' OR '1'='1"
         escaped = escape_sql_string(input_str)
@@ -171,23 +171,23 @@ class TestEscapeSQLString:
 class TestEscapeIdentifier:
     """Test escape_identifier function."""
 
-    def test_escape_uuid(self):
+    def test_escape_uuid(self) -> None:
         """Test escaping UUID identifiers."""
         uuid = "550e8400-e29b-41d4-a716-446655440000"
         assert escape_identifier(uuid) == uuid
 
-    def test_escape_alphanumeric_id(self):
+    def test_escape_alphanumeric_id(self) -> None:
         """Test escaping alphanumeric identifiers."""
         assert escape_identifier("tenant_123") == "tenant_123"
         assert escape_identifier("app-abc-123") == "app-abc-123"
 
-    def test_escape_identifier_with_quote(self):
+    def test_escape_identifier_with_quote(self) -> None:
         """Test escaping identifier with single quote."""
         malicious = "tenant' OR '1'='1"
         escaped = escape_identifier(malicious)
         assert escaped == "tenant'' OR ''1''=''1"
 
-    def test_identifier_injection_attempt(self):
+    def test_identifier_injection_attempt(self) -> None:
         """Test prevention of injection through identifiers."""
         # Common identifier injection patterns
         test_cases = [
@@ -203,7 +203,7 @@ class TestEscapeIdentifier:
 class TestSQLInjectionIntegration:
     """Integration tests simulating real SQL construction scenarios."""
 
-    def test_complete_where_clause_safety(self):
+    def test_complete_where_clause_safety(self) -> None:
         """Test that a complete WHERE clause is safe from injection."""
         # Simulating typical query construction
         tenant_id = "tenant' OR '1'='1"
@@ -229,7 +229,7 @@ class TestSQLInjectionIntegration:
         # Verify SQL structure is preserved (3 conditions with AND)
         assert sql.count("AND") == 2
 
-    def test_multiple_conditions_with_injection_attempts(self):
+    def test_multiple_conditions_with_injection_attempts(self) -> None:
         """Test multiple conditions all attempting injection."""
         conditions = {
             "tenant_id": "t1' OR tenant_id='t2",
@@ -265,7 +265,7 @@ class TestSQLInjectionIntegration:
             ("' OR 1=1 LIMIT 1 --", "Limit bypass"),
         ],
     )
-    def test_common_injection_vectors(self, attack_vector, description):
+    def test_common_injection_vectors(self, attack_vector: str, description: str) -> None:
         """Test protection against common injection attack vectors."""
         escaped = escape_sql_string(attack_vector)
 
@@ -283,7 +283,7 @@ class TestSQLInjectionIntegration:
         # Verify SQL has exactly 2 quotes (opening and closing)
         assert sql.count("'") >= 2  # At least opening and closing
 
-    def test_logstore_specific_scenario(self):
+    def test_logstore_specific_scenario(self) -> None:
         """Test SQL injection prevention in LogStore-specific scenarios."""
         # Simulate LogStore query with window function
         tenant_id = "tenant' OR '1'='1"
@@ -319,7 +319,7 @@ class TestSQLInjectionIntegration:
 class TestLogStoreQueryEscape:
     """Test escape_logstore_query_value for SDK mode query syntax."""
 
-    def test_normal_value(self):
+    def test_normal_value(self) -> None:
         """Test escaping normal alphanumeric value."""
         value = "550e8400-e29b-41d4-a716-446655440000"
         escaped = escape_logstore_query_value(value)
@@ -327,11 +327,11 @@ class TestLogStoreQueryEscape:
         # Should be wrapped in double quotes
         assert escaped == '"550e8400-e29b-41d4-a716-446655440000"'
 
-    def test_empty_value(self):
+    def test_empty_value(self) -> None:
         """Test escaping empty string."""
         assert escape_logstore_query_value("") == '""'
 
-    def test_value_with_and_keyword(self):
+    def test_value_with_and_keyword(self) -> None:
         """Test that 'and' keyword is neutralized when quoted."""
         malicious = "value and field:evil"
         escaped = escape_logstore_query_value(malicious)
@@ -343,7 +343,7 @@ class TestLogStoreQueryEscape:
         query = f"tenant_id:{escaped}"
         assert query == 'tenant_id:"value and field:evil"'
 
-    def test_value_with_or_keyword(self):
+    def test_value_with_or_keyword(self) -> None:
         """Test that 'or' keyword is neutralized when quoted."""
         malicious = "tenant_a or tenant_id:tenant_b"
         escaped = escape_logstore_query_value(malicious)
@@ -353,14 +353,14 @@ class TestLogStoreQueryEscape:
         query = f"tenant_id:{escaped}"
         assert "or" in query  # Present but as literal string
 
-    def test_value_with_not_keyword(self):
+    def test_value_with_not_keyword(self) -> None:
         """Test that 'not' keyword is neutralized when quoted."""
         malicious = "not field:value"
         escaped = escape_logstore_query_value(malicious)
 
         assert escaped == '"not field:value"'
 
-    def test_value_with_parentheses(self):
+    def test_value_with_parentheses(self) -> None:
         """Test that parentheses are neutralized when quoted."""
         malicious = "(tenant_a or tenant_b)"
         escaped = escape_logstore_query_value(malicious)
@@ -369,7 +369,7 @@ class TestLogStoreQueryEscape:
         assert "(" in escaped  # Present as literal
         assert ")" in escaped  # Present as literal
 
-    def test_value_with_colon(self):
+    def test_value_with_colon(self) -> None:
         """Test that colons are neutralized when quoted."""
         malicious = "field:value"
         escaped = escape_logstore_query_value(malicious)
@@ -377,7 +377,7 @@ class TestLogStoreQueryEscape:
         assert escaped == '"field:value"'
         assert ":" in escaped  # Present as literal
 
-    def test_value_with_double_quotes(self):
+    def test_value_with_double_quotes(self) -> None:
         """Test that internal double quotes are escaped."""
         value_with_quotes = 'tenant"test"value'
         escaped = escape_logstore_query_value(value_with_quotes)
@@ -387,7 +387,7 @@ class TestLogStoreQueryEscape:
         # Should have outer quotes plus escaped inner quotes
         assert '\\"' in escaped
 
-    def test_value_with_backslash(self):
+    def test_value_with_backslash(self) -> None:
         """Test that backslashes are escaped."""
         value_with_backslash = "tenant\\test"
         escaped = escape_logstore_query_value(value_with_backslash)
@@ -396,7 +396,7 @@ class TestLogStoreQueryEscape:
         assert escaped == '"tenant\\\\test"'
         assert "\\\\" in escaped
 
-    def test_value_with_backslash_and_quote(self):
+    def test_value_with_backslash_and_quote(self) -> None:
         """Test escaping both backslash and double quote."""
         value = 'path\\to\\"file"'
         escaped = escape_logstore_query_value(value)
@@ -407,7 +407,7 @@ class TestLogStoreQueryEscape:
         assert "\\\\" in escaped  # Escaped backslash
         assert '\\"' in escaped  # Escaped double quote
 
-    def test_complex_injection_attempt(self):
+    def test_complex_injection_attempt(self) -> None:
         """Test complex injection combining multiple operators."""
         malicious = 'tenant_a" or (tenant_id:"tenant_b" and app_id:"evil")'
         escaped = escape_logstore_query_value(malicious)
@@ -420,19 +420,19 @@ class TestLogStoreQueryEscape:
         assert "and" in escaped
         assert '\\"' in escaped  # Escaped quotes
 
-    def test_only_backslash(self):
+    def test_only_backslash(self) -> None:
         """Test escaping a single backslash."""
         assert escape_logstore_query_value("\\") == '"\\\\"'
 
-    def test_only_double_quote(self):
+    def test_only_double_quote(self) -> None:
         """Test escaping a single double quote."""
         assert escape_logstore_query_value('"') == '"\\""'
 
-    def test_multiple_backslashes(self):
+    def test_multiple_backslashes(self) -> None:
         """Test escaping multiple consecutive backslashes."""
         assert escape_logstore_query_value("\\\\\\") == '"\\\\\\\\\\\\"'  # 3 backslashes -> 6
 
-    def test_escape_sequence_like_input(self):
+    def test_escape_sequence_like_input(self) -> None:
         """Test that existing escape sequences are properly escaped."""
         # Input looks like already escaped, but we still escape it
         value = 'value\\"test'
@@ -454,7 +454,7 @@ class TestLogStoreQueryEscape:
         ("Backslash escape bypass", "app_id", "app\\ and app_id:evil"),
     ],
 )
-def test_logstore_query_injection_scenarios(attack_scenario: str, field: str, malicious_value: str):
+def test_logstore_query_injection_scenarios(attack_scenario: str, field: str, malicious_value: str) -> None:
     """Test that various LogStore query injection attempts are neutralized."""
     escaped = escape_logstore_query_value(malicious_value)
 

api/tests/unit_tests/extensions/otel/decorators/handlers/test_generate_handler.py

@@ -6,11 +6,14 @@ Test objectives:
 2. Verify span attribute mapping correctness
 """
 
-from unittest.mock import MagicMock
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+from sqlalchemy.orm import Session
 
 from core.app.entities.app_invoke_entities import InvokeFrom
 from extensions.otel.decorators.handlers.generate_handler import AppGenerateHandler
 from extensions.otel.semconv import DifySpanAttributes, GenAIAttributes
+from models import Account, App
 from tests.helpers.unit.config_override import config_overrides_context
 
 
@@ -19,8 +22,12 @@ class TestAppGenerateHandler:
 
     @config_overrides_context(ENABLE_OTEL=True)
     def test_compatible_with_real_function_signature(
-        self, tracer_provider_with_memory_exporter, mock_app_model, mock_account_user
-    ):
+        self,
+        tracer_provider_with_memory_exporter: TracerProvider,
+        mock_app_model: App,
+        mock_account_user: Account,
+        sqlite_session: Session,
+    ) -> None:
         """
         Verify handler compatibility with real AppGenerateService.generate signature.
 
@@ -31,17 +38,16 @@ class TestAppGenerateHandler:
 
         handler = AppGenerateHandler()
 
-        kwargs = {
-            "session": MagicMock(),
-            "app_model": mock_app_model,
-            "user": mock_account_user,
-            "args": {"workflow_id": "test-wf-123"},
-            "invoke_from": InvokeFrom.DEBUGGER,
-            "streaming": True,
-            "root_node_id": None,
-        }
-
-        arguments = handler._extract_arguments(AppGenerateService.generate, **kwargs)
+        arguments = handler._extract_arguments(
+            AppGenerateService.generate,
+            session=sqlite_session,
+            app_model=mock_app_model,
+            user=mock_account_user,
+            args={"workflow_id": "test-wf-123"},
+            invoke_from=InvokeFrom.DEBUGGER,
+            streaming=True,
+            root_node_id=None,
+        )
 
         assert arguments is not None, "Failed to extract arguments from AppGenerateService.generate"
         assert "app_model" in arguments, "Handler uses app_model but parameter is missing"
@@ -51,8 +57,12 @@ class TestAppGenerateHandler:
 
     @config_overrides_context(ENABLE_OTEL=True)
     def test_all_span_attributes_set_correctly(
-        self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_app_model, mock_account_user
-    ):
+        self,
+        tracer_provider_with_memory_exporter: TracerProvider,
+        memory_span_exporter: InMemorySpanExporter,
+        mock_app_model: App,
+        mock_account_user: Account,
+    ) -> None:
         """Verify all span attributes are mapped correctly"""
         handler = AppGenerateHandler()
         tracer = tracer_provider_with_memory_exporter.get_tracer(__name__)
@@ -66,7 +76,9 @@ class TestAppGenerateHandler:
         mock_app_model.tenant_id = test_tenant_id
         mock_account_user.id = test_user_id
 
-        def dummy_func(app_model, user, args, invoke_from, streaming=True):
+        def dummy_func(
+            app_model: App, user: Account, args: dict[str, object], invoke_from: InvokeFrom, streaming: bool = True
+        ) -> str:
             return "result"
 
         handler.wrapper(
@@ -82,6 +94,7 @@ class TestAppGenerateHandler:
         spans = memory_span_exporter.get_finished_spans()
         assert len(spans) == 1
         attrs = spans[0].attributes
+        assert attrs is not None
 
         assert attrs[DifySpanAttributes.APP_ID] == test_app_id
         assert attrs[DifySpanAttributes.TENANT_ID] == test_tenant_id

api/tests/unit_tests/extensions/otel/decorators/handlers/test_workflow_app_runner_handler.py

@@ -6,6 +6,8 @@ Test objectives:
 2. Verify span attribute mapping correctness
 """
 
+from unittest.mock import MagicMock
+
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
@@ -41,7 +43,7 @@ def test_service_driver_emits_workflow_span(sqlite_engine: Engine) -> None:
 class TestWorkflowAppRunnerHandler:
     """Core tests for WorkflowAppRunnerHandler"""
 
-    def test_handler_structure_dependencies(self):
+    def test_handler_structure_dependencies(self) -> None:
         """
         Verify handler dependencies on WorkflowAppRunner structure.
 
@@ -69,8 +71,11 @@ class TestWorkflowAppRunnerHandler:
 
     @config_overrides_context(ENABLE_OTEL=True)
     def test_all_span_attributes_set_correctly(
-        self, tracer_provider_with_memory_exporter, memory_span_exporter, mock_workflow_runner
-    ):
+        self,
+        tracer_provider_with_memory_exporter: TracerProvider,
+        memory_span_exporter: InMemorySpanExporter,
+        mock_workflow_runner: MagicMock,
+    ) -> None:
         """Verify all span attributes are mapped correctly"""
         handler = WorkflowAppRunnerHandler()
         tracer = tracer_provider_with_memory_exporter.get_tracer(__name__)
@@ -86,7 +91,7 @@ class TestWorkflowAppRunnerHandler:
         mock_workflow_runner.application_generate_entity.app_config.tenant_id = test_tenant_id
         mock_workflow_runner.application_generate_entity.app_config.workflow_id = test_workflow_id
 
-        def runner_run(runner):
+        def runner_run(runner: MagicMock) -> str:
             return "result"
 
         handler.wrapper(tracer, runner_run, mock_workflow_runner)
@@ -94,6 +99,7 @@ class TestWorkflowAppRunnerHandler:
         spans = memory_span_exporter.get_finished_spans()
         assert len(spans) == 1
         attrs = spans[0].attributes
+        assert attrs is not None
 
         assert attrs[DifySpanAttributes.APP_ID] == test_app_id
         assert attrs[DifySpanAttributes.TENANT_ID] == test_tenant_id

api/tests/unit_tests/extensions/test_ext_login.py

@@ -1,5 +1,5 @@
 import json
-from collections.abc import Callable
+from collections.abc import Callable, Generator
 from typing import cast
 from unittest import mock
 
@@ -12,10 +12,11 @@ from core.logging.context import clear_request_context, get_identity_context
 from extensions import ext_login
 from extensions.ext_login import unauthorized_handler
 from models.account import TenantAccountRole
+from models.enums import EndUserType
 
 
 @pytest.fixture(autouse=True)
-def _reset_logging_context():
+def _reset_logging_context() -> Generator[None]:
     clear_request_context()
     yield
     clear_request_context()
@@ -50,7 +51,7 @@ def test_on_user_logged_in_sets_end_user_logging_identity() -> None:
     end_user = ext_login.EndUser(
         id="end-user-id",
         tenant_id="tenant-id",
-        type="browser",
+        type=EndUserType.BROWSER,
     )
     clear_request_context()