diff --git a/api/tests/unit_tests/events/test___init__.py b/api/tests/unit_tests/events/test___init__.py index 1837bfa6b32..44aff41fa21 100644 --- a/api/tests/unit_tests/events/test___init__.py +++ b/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 diff --git a/api/tests/unit_tests/events/test_app_event.py b/api/tests/unit_tests/events/test_app_event.py index 7519487a808..34a428d7516 100644 --- a/api/tests/unit_tests/events/test_app_event.py +++ b/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 diff --git a/api/tests/unit_tests/extensions/logstore/test_sql_escape.py b/api/tests/unit_tests/extensions/logstore/test_sql_escape.py index 63172b3f9b7..8a7baf7654e 100644 --- a/api/tests/unit_tests/extensions/logstore/test_sql_escape.py +++ b/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) diff --git a/api/tests/unit_tests/extensions/otel/decorators/handlers/test_generate_handler.py b/api/tests/unit_tests/extensions/otel/decorators/handlers/test_generate_handler.py index e5ccc0aca6b..93c068079f7 100644 --- a/api/tests/unit_tests/extensions/otel/decorators/handlers/test_generate_handler.py +++ b/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 diff --git a/api/tests/unit_tests/extensions/otel/decorators/handlers/test_workflow_app_runner_handler.py b/api/tests/unit_tests/extensions/otel/decorators/handlers/test_workflow_app_runner_handler.py index 54e54fe0dea..5b67283abcc 100644 --- a/api/tests/unit_tests/extensions/otel/decorators/handlers/test_workflow_app_runner_handler.py +++ b/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 diff --git a/api/tests/unit_tests/extensions/test_ext_login.py b/api/tests/unit_tests/extensions/test_ext_login.py index 345a00e8e7c..fb1db799bd6 100644 --- a/api/tests/unit_tests/extensions/test_ext_login.py +++ b/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()