dabla commented on code in PR #74074: URL: https://github.com/apache/airflow/pull/74074#discussion_r4562046825
########## airflow-core/tests/unit/core/test_process_context.py: ########## @@ -0,0 +1,61 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import sys +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from airflow.process_context import force_server_context, should_use_task_sdk_api_path + +TASK_RUNNER = "airflow.sdk.execution_time.task_runner" + + +class TestShouldUseTaskSdkApiPath: + """Whether models.Variable/Connection take the Execution API path, per thread.""" + + def test_false_when_the_task_runner_was_never_imported(self, monkeypatch): Review Comment: Right, neither touches the changed lines. Both dropped in 82a0f4b15e. ########## airflow-core/tests/unit/models/test_taskinstance.py: ########## @@ -5668,3 +5668,86 @@ def test_mapped_archival_ack_uses_fork_facts_when_definition_is_missing(loop_run archived = session.get(TaskInstance, archiving_id) assert (archived.working_set, archived.archived_reason) == (None, "superseded") + + +class TestInProcessSupervisorFromThreads: + """ + Under the in-process supervisor (dag.test(), dag_maker) a task's own threads make SDK calls + while others are being served: every call must go through and get its own answer. + """ + + @staticmethod + def _run(dag_maker, session): + dag_run = dag_maker.create_dagrun() + (ti,) = dag_run.task_instance_scheduling_decisions(session=session).schedulable_tis + with contextlib.suppress(BaseException): Review Comment: Dropped in 82a0f4b15e. The tests now call `dag.test()` and assert on the run's state, see the thread below. ########## airflow-core/tests/unit/models/test_taskinstance.py: ########## @@ -5668,3 +5668,86 @@ def test_mapped_archival_ack_uses_fork_facts_when_definition_is_missing(loop_run archived = session.get(TaskInstance, archiving_id) assert (archived.working_set, archived.archived_reason) == (None, "superseded") + + +class TestInProcessSupervisorFromThreads: Review Comment: Moved to `test_dag.py`, after `test_dag_test_with_dependencies`, and they call `dag.test()` itself now instead of `dag_maker.run_ti` (82a0f4b15e). You are right that on main they only fail when the threads overlap; they are the end-to-end check. The deterministic one is `test_only_the_code_serving_a_request_stops_seeing_the_comms` in the Task SDK, which forces the overlap with two events. ########## task-sdk/tests/task_sdk/execution_time/test_supervisor.py: ########## @@ -5319,6 +5320,91 @@ def fake_api_client(dag, attempt_id): assert isinstance(collected[0], _Failure) +class TestInProcessSupervisorCommsAcrossThreads: + """ + dag.test() serves a task's requests in the task's own process, possibly for several of the task's + threads at once. While a request is served, the code serving it stops seeing the comms, and only + that code: the task's other threads keep theirs. The in-process API server's side is covered in the + execution API's test_app.py. + """ + + @staticmethod + def _comms(serve): + """In-process comms whose supervisor answers each request with ``serve(msg)``.""" + comms = InProcessSupervisorComms(supervisor=mock.Mock(spec=InProcessTestSupervisor)) + + def handle_request(msg, log, req_id): + sink = _IN_PROCESS_RESPONSE_SINK.get() + assert sink is not None + sink.append(serve(msg)) + + comms.supervisor._handle_request.side_effect = handle_request + return comms + + def test_only_the_code_serving_a_request_stops_seeing_the_comms(self): + """ + While one thread's request is served, the task's other threads must keep their comms: seeing + none, they would take their Variable and Connection lookups for the server side and find nothing. + """ + import threading Review Comment: Removed in 82a0f4b15e, together with the local `task_runner` imports in the same class, which line 82 already has. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
