This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 5bc8796e3a0 Run standard venv operator tests DB-free so they
parallelize under xdist (#68533)
5bc8796e3a0 is described below
commit 5bc8796e3a0d2ed3bb7cf9a45bb80a5588c2607d
Author: Jarek Potiuk <[email protected]>
AuthorDate: Sun Jun 21 11:18:32 2026 -0400
Run standard venv operator tests DB-free so they parallelize under xdist
(#68533)
* Run standard venv operator tests DB-free so they parallelize under xdist
PythonVirtualenvOperator/ExternalPythonOperator tests dominate the
`Providers[standard]` suite (~450s of ~510s) because they build real
virtualenvs, and they were forced to run serially: marked `db_test`,
so pytest-xdist could not parallelize them.
Their only real DB coupling was the test harness, not the metadata DB —
but the venv subprocess reconnects to the supervisor over a socket, so
they cannot use the plain `run_task` mock (no real socket → `OSError:
Socket operation on non-socket`).
Add an optional `client=` parameter to `InProcessTestSupervisor.start()`
(and `run_task_in_process`) so a task can run through the real supervisor
socketpair with an injected Execution-API client instead of the
DB-backed in-process API server.
Add `tests_common.test_utils.in_process_taskrun.run_task_no_db`, which
uses a real `Client(dry_run=True)` whose transport remembers XCom writes
in an in-memory dict — so a venv operator runs with a working supervisor
socket and no metadata DB, and its pushed XCom can be asserted.
Convert `TestPythonVirtualenvOperator` and `TestExternalPythonOperator`
onto it (via a `_DBFreeVenvRun` mixin) and drop their `db_test` mark, so
they parallelize under `--skip-db-tests --use-xdist`. On Airflow 2.x (no
`InProcessTestSupervisor`) the classes fall back to the DB-backed path
and keep `db_test`; on Airflow 3.x releases without the new `client=`
param the helper falls back to overriding `_api_client`. Branch venv
classes keep `db_test` (their multi-task assertions need a real DAG run).
Measured: the two classes run ~3x faster at `-n 4` (157 tests, 89s vs
~266s serial); larger gains at CI worker counts.
* Drop dead b"null" fallback in DB-free XCom write parsing
A POST to /xcoms/... always carries the serialized value (even None
serializes to b"null"), so request.content is never empty and the
`or b"null"` fallback was dead. Addresses review feedback.
---
.../tests_common/test_utils/in_process_taskrun.py | 145 +++++++++++++++++++
.../tests/unit/standard/operators/test_python.py | 153 +++++++++++++++++++--
.../src/airflow/sdk/execution_time/supervisor.py | 17 ++-
3 files changed, 301 insertions(+), 14 deletions(-)
diff --git a/devel-common/src/tests_common/test_utils/in_process_taskrun.py
b/devel-common/src/tests_common/test_utils/in_process_taskrun.py
new file mode 100644
index 00000000000..1dd3f1c736a
--- /dev/null
+++ b/devel-common/src/tests_common/test_utils/in_process_taskrun.py
@@ -0,0 +1,145 @@
+# 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.
+"""DB-free, xdist-safe execution of a task through a *real* supervisor socket.
+
+`run_task` (in ``pytest_plugin``) mocks supervisor comms entirely in-process
and
+has **no real socket**, so operators that spawn a subprocess which re-connects
to
+the supervisor — ``PythonVirtualenvOperator``, ``ExternalPythonOperator``,
+``run_as_user`` — fail there with ``OSError: Socket operation on non-socket``.
+
+This drives the *real* ``InProcessTestSupervisor`` (its socketpair machinery is
+created specifically for VirtualEnv operators) but injects a **dry-run
Execution-API
+client** instead of the DB-backed in-process API server, so the subprocess
gets a
+working supervisor socket without touching the metadata DB. Tests using it
need no
+``@pytest.mark.db_test`` and run under xdist.
+
+The client is the real ``Client(dry_run=True)`` (which already fakes the run
+context and no-ops heartbeats via ``noop_handler``), with the discarding
transport
+swapped for one that *remembers* XCom writes in an in-memory dict — exposed as
+``client.pushed_xcoms`` so tests can assert what a task pushed.
+
+Requires the Task SDK ``run_task_in_process(..., client=)`` parameter (newer
than
+Airflow 3.0/3.1, and absent on 2.x); callers must gate on its availability and
fall
+back to a DB-backed path otherwise.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+ from collections.abc import Callable
+
+ from airflow.sdk.api.client import Client
+ from airflow.sdk.execution_time.supervisor import TaskRunResult
+ from airflow.sdk.types import Operator
+
+# XCom is the only resource that must round-trip; the run-context is fed back
from the
+# (valid) ti_context the test built, and everything else (heartbeat, state
updates) is the
+# stock ``noop_handler``. (``noop_handler``'s own run-context is stale vs the
live schema.)
+_XCOM_PATH_PARTS = 5 # /xcoms/{dag_id}/{run_id}/{task_id}/{key}
+
+
+def _remembering_handler(store: dict, run_context_json: bytes) -> Callable:
+ """A dry-run transport handler: valid run-context + XCom round-trip from
``store``, else no-op."""
+ import httpx
+
+ from airflow.sdk.api.client import noop_handler
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ path = request.url.path
+ if path.startswith("/task-instances/") and path.endswith("/run"):
+ return httpx.Response(200, content=run_context_json)
+ parts = path.strip("/").split("/")
+ if len(parts) == _XCOM_PATH_PARTS and parts[0] == "xcoms":
+ dag_id, run_id, task_id, key = parts[1:]
+ sig = (dag_id, run_id, task_id, key)
+ if request.method == "POST":
+ store[sig] = json.loads(request.content)
+ return httpx.Response(201, json={"ok": True})
+ if request.method == "GET":
+ if sig in store:
+ return httpx.Response(200, json={"key": key, "value":
store[sig]})
+ return httpx.Response(404, json={"detail": "XCom not found"})
+ return noop_handler(request)
+
+ return handler
+
+
+def build_in_memory_client(ti_context) -> Client:
+ """A real ``Client(dry_run=True)`` that remembers XCom writes (no DB, no
network).
+
+ ``ti_context`` (a ``TIRunContext``) is replayed for the task-start
request. Pushed XCom
+ values are exposed as ``client.pushed_xcoms`` keyed by ``(dag_id, run_id,
task_id, key)``.
+ """
+ import httpx
+
+ from airflow.sdk.api.client import Client
+
+ store: dict[tuple[str, str, str, str], Any] = {}
+ client = Client(
+ base_url=None,
+ dry_run=True,
+ token="",
+ transport=httpx.MockTransport(_remembering_handler(store,
ti_context.model_dump_json().encode())),
+ )
+ client.pushed_xcoms = store # type: ignore[attr-defined]
+ return client
+
+
+def pushed_xcom(xcoms: dict, ti, key: str = "return_value") -> Any:
+ """Read an XCom a task pushed during :func:`run_task_no_db` (``None`` if
absent)."""
+ return xcoms.get((ti.dag_id, ti.run_id, ti.task_id, key))
+
+
+def run_task_no_db(
+ task: Operator,
+ create_runtime_ti: Callable[..., Any],
+ *,
+ logical_date: Any | None = None,
+) -> tuple[TaskRunResult, dict]:
+ """Run *task* DB-free through the real-socket in-process supervisor.
+
+ Returns ``(result, pushed_xcoms)`` where ``result`` is the stock
``TaskRunResult``
+ (``.state`` / ``.msg`` / ``.error`` / ``.ti``) and ``pushed_xcoms`` is the
dict of
+ XCom values the task pushed (read via :func:`pushed_xcom`).
+ """
+ from uuid6 import uuid7
+
+ from airflow.sdk.api.datamodels._generated import TaskInstance as
TaskInstanceDTO
+ from airflow.sdk.execution_time.supervisor import run_task_in_process
+
+ ti_kwargs = {} if logical_date is None else {"logical_date": logical_date}
+ rti = create_runtime_ti(task, **ti_kwargs)
+
+ # `start()` model_dumps `what`; the plain DTO dumps cleanly, whereas the
+ # operator-laden RuntimeTaskInstance trips forward refs
(RetryPolicy/WeightRuleParam).
+ what = TaskInstanceDTO(
+ id=rti.id,
+ task_id=rti.task_id,
+ dag_id=rti.dag_id,
+ run_id=rti.run_id,
+ try_number=rti.try_number,
+ map_index=rti.map_index,
+ dag_version_id=uuid7(),
+ queue="default",
+ )
+
+ client = build_in_memory_client(rti._ti_context_from_server)
+ result = run_task_in_process(what, task, client=client)
+ return result, client.pushed_xcoms # type: ignore[attr-defined]
diff --git a/providers/standard/tests/unit/standard/operators/test_python.py
b/providers/standard/tests/unit/standard/operators/test_python.py
index f432ec8b5ee..d8228b1032f 100644
--- a/providers/standard/tests/unit/standard/operators/test_python.py
+++ b/providers/standard/tests/unit/standard/operators/test_python.py
@@ -23,6 +23,7 @@ import logging
import os
import pickle
import re
+import shutil
import sys
import tempfile
import warnings
@@ -65,6 +66,7 @@ from airflow.utils.types import DagRunType
from tests_common.test_utils.compat import TriggerRule, timezone
from tests_common.test_utils.db import clear_db_runs
+from tests_common.test_utils.in_process_taskrun import pushed_xcom,
run_task_no_db
from tests_common.test_utils.taskinstance import get_template_context,
run_task_instance
from tests_common.test_utils.version_compat import (
AIRFLOW_V_3_0_1,
@@ -86,7 +88,6 @@ if TYPE_CHECKING:
from airflow.models.dagrun import DagRun
from airflow.sdk import Context
-pytestmark = [pytest.mark.db_test, pytest.mark.need_serialized_dag]
AIRFLOW_ROOT_PATH = Path(__file__).parents[6]
@@ -202,12 +203,21 @@ class BasePythonTest:
return ti
return ti.task
+ @staticmethod
+ def _pull_xcom(ran):
+ """Pull the return-value XCom for whatever
``run_as_task(return_ti=True)`` returned.
+
+ Overridden in the DB-free venv mixin to read from the in-memory
runtime result.
+ """
+ return TaskInstance.xcom_pull(ran)
+
def render_templates(self, fn, **kwargs):
"""Create TaskInstance and render templates without actual run."""
return self.create_ti(fn, **kwargs).render_templates()
class TestPythonOperator(BasePythonTest):
+ pytestmark = [pytest.mark.db_test, pytest.mark.need_serialized_dag]
opcls = PythonOperator
@pytest.fixture(autouse=True)
@@ -400,6 +410,7 @@ class TestPythonOperator(BasePythonTest):
class TestBranchOperator(BasePythonTest):
+ pytestmark = [pytest.mark.db_test, pytest.mark.need_serialized_dag]
opcls = BranchPythonOperator
@pytest.fixture(autouse=True)
@@ -618,6 +629,7 @@ class TestBranchOperator(BasePythonTest):
class TestShortCircuitOperator(BasePythonTest):
+ pytestmark = [pytest.mark.db_test, pytest.mark.need_serialized_dag]
opcls = ShortCircuitOperator
@pytest.fixture(autouse=True)
@@ -952,6 +964,8 @@ class TestDagBundleImportInSubprocess(BasePythonTest):
from their Dag bundle by verifying PYTHONPATH is correctly set (Airflow
3.x+).
"""
+ pytestmark = [pytest.mark.db_test, pytest.mark.need_serialized_dag]
+
@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Dag Bundle import fix
is for Airflow 3.x+")
@mock.patch("airflow.providers.standard.operators.python._execute_in_subprocess")
def test_dag_bundle_import_in_subprocess(
@@ -1046,7 +1060,7 @@ class BaseTestPythonVirtualenvOperator(BasePythonTest):
return None
ti = self.run_as_task(f, return_ti=True)
- assert TaskInstance.xcom_pull(ti) is None
+ assert self._pull_xcom(ti) is None
def test_return_false(self):
def f():
@@ -1054,7 +1068,7 @@ class BaseTestPythonVirtualenvOperator(BasePythonTest):
ti = self.run_as_task(f, return_ti=True)
- assert TaskInstance.xcom_pull(ti) is False
+ assert self._pull_xcom(ti) is False
def test_lambda(self):
with pytest.raises(
@@ -1072,8 +1086,9 @@ class BaseTestPythonVirtualenvOperator(BasePythonTest):
def f(templates_dict):
return templates_dict["ds"]
- task = self.run_as_task(f, templates_dict={"ds": "{{ ds }}"})
- assert task.templates_dict == {"ds": self.ds_templated}
+ # the callable receives (and returns) the rendered templates_dict value
+ ti = self.run_as_task(f, return_ti=True, templates_dict={"ds": "{{ ds
}}"})
+ assert self._pull_xcom(ti) == self.ds_templated
@pytest.mark.parametrize(
"serializer",
@@ -1112,6 +1127,8 @@ class BaseTestPythonVirtualenvOperator(BasePythonTest):
with pytest.raises(AirflowException, match=r"cannot be
pickled.*\['bad_obj'\]"):
op._write_args(tmp_path / "args.pkl")
+ @pytest.mark.db_test
+ @pytest.mark.need_serialized_dag
def test_virtualenv_serializable_context_fields(self,
create_task_instance):
"""Ensure all template context fields are listed in the operator.
@@ -1244,7 +1261,7 @@ class BaseTestPythonVirtualenvOperator(BasePythonTest):
return os.environ["MY_ENV_VAR"]
ti = self.run_as_task(f, env_vars={"MY_ENV_VAR": "ABCDE"},
return_ti=True)
- assert TaskInstance.xcom_pull(ti) == "ABCDE"
+ assert self._pull_xcom(ti) == "ABCDE"
def test_environment_variables_with_inherit_env_true(self, monkeypatch):
monkeypatch.setenv("MY_ENV_VAR", "QWERT")
@@ -1255,7 +1272,7 @@ class BaseTestPythonVirtualenvOperator(BasePythonTest):
return os.environ["MY_ENV_VAR"]
ti = self.run_as_task(f, inherit_env=True, return_ti=True)
- assert TaskInstance.xcom_pull(ti) == "QWERT"
+ assert self._pull_xcom(ti) == "QWERT"
def test_environment_variables_with_inherit_env_false(self, monkeypatch):
monkeypatch.setenv("MY_ENV_VAR", "TYUIO")
@@ -1277,18 +1294,123 @@ class BaseTestPythonVirtualenvOperator(BasePythonTest):
return os.environ["MY_ENV_VAR"]
ti = self.run_as_task(f, env_vars={"MY_ENV_VAR": "EFGHI"},
inherit_env=True, return_ti=True)
- assert TaskInstance.xcom_pull(ti) == "EFGHI"
+ assert self._pull_xcom(ti) == "EFGHI"
venv_cache_path = tempfile.mkdtemp(prefix="venv_cache_path")
[email protected](scope="session", autouse=True)
+def _cleanup_venv_cache_path():
+ """Remove the shared (per-process / per-xdist-worker) venv cache dir after
the session."""
+ yield
+ shutil.rmtree(venv_cache_path, ignore_errors=True)
+
+
+class _DBFreeVenvRun:
+ """Run venv-operator tests DB-free via the real-socket in-process
supervisor.
+
+ PythonVirtualenvOperator / ExternalPythonOperator spawn a subprocess that
+ reconnects to the supervisor over a socket, so they cannot run under the
+ plain ``run_task`` mock (no socket). This mixin overrides
``BasePythonTest``'s
+ DB-backed execution to use ``run_task_no_db`` (real socketpair + in-memory
+ backend), so these classes need no ``db_test`` mark and run under xdist.
+ """
+
+ @pytest.fixture(autouse=True)
+ def base_tests_setup(self, request, create_runtime_ti): # overrides the
DB-backed base fixture
+ from airflow.sdk import DAG
+
+ self.dag_id = f"dag_{slugify(request.cls.__name__)}"
+ self.task_id = f"task_{slugify(request.node.name, max_length=40)}"
+ self.run_id = f"run_{slugify(request.node.name, max_length=40)}"
+ self.ds_templated = self.default_date.date().isoformat()
+ self._create_runtime_ti = create_runtime_ti
+ # A plain (un-serialized, no-DB) DAG for tests that instantiate
operators directly.
+ self.dag_non_serialized = DAG(
+ self.dag_id, schedule=None, start_date=self.default_date,
template_searchpath=TEMPLATE_SEARCHPATH
+ )
+
+ def _run_dbfree(self, fn, **kwargs):
+ from airflow.sdk import DAG
+
+ # Fresh DAG per run, carrying template_searchpath so templated fields
+ # (e.g. requirements="requirements.txt") resolve like the DB path's
dag_maker.
+ dag = DAG(
+ self.dag_id,
+ schedule=None,
+ start_date=self.default_date,
+ template_searchpath=TEMPLATE_SEARCHPATH,
+ )
+ with dag:
+ task = self.opcls(task_id=self.task_id, python_callable=fn,
**self.default_kwargs(**kwargs))
+ # keep the pushed XComs so assertions can read them back
+ result, self._last_xcoms = run_task_no_db(
+ task, self._create_runtime_ti, logical_date=self.default_date
+ )
+ return task, result
+
+ def run_as_task(self, fn, return_ti=False, **kwargs):
+ task, result = self._run_dbfree(fn, **kwargs)
+ if return_ti:
+ return result
+ if result.error is not None:
+ raise result.error
+ return task
+
+ def run_as_operator(self, fn, **kwargs):
+ task, result = self._run_dbfree(fn, **kwargs)
+ if result.error is not None:
+ raise result.error
+ return task
+
+ def _pull_xcom(self, ran): # ``ran`` is the TaskRunResult from
run_as_task(return_ti=True)
+ return pushed_xcom(self._last_xcoms, ran.ti)
+
+
+def _dbfree_venv_supported() -> bool:
+ """Whether venv operators can run DB-free here.
+
+ Needs the Task SDK ``InProcessTestSupervisor`` with the ``client=``
parameter on
+ ``run_task_in_process`` (no ``InProcessTestSupervisor`` on Airflow 2.x;
the parameter
+ is newer than 3.0/3.1, where the dry-run client cannot serve all calls).
+ """
+ if not AIRFLOW_V_3_0_PLUS:
+ return False
+ try:
+ import inspect
+
+ from airflow.sdk.execution_time.supervisor import run_task_in_process
+ except ImportError:
+ return False
+ return "client" in inspect.signature(run_task_in_process).parameters
+
+
+_DBFREE_VENV = _dbfree_venv_supported()
+
+# Where supported, run the venv classes DB-free (no ``db_test``,
xdist-friendly) via the
+# ``_DBFreeVenvRun`` mixin; otherwise fall back to the DB-backed
BasePythonTest path.
+if _DBFREE_VENV:
+
+ class _VenvTestBase(_DBFreeVenvRun, BaseTestPythonVirtualenvOperator):
+ pass
+
+ _VENV_DB_MARKS: list = []
+else:
+
+ class _VenvTestBase(BaseTestPythonVirtualenvOperator): # type:
ignore[no-redef]
+ pass
+
+ _VENV_DB_MARKS = [pytest.mark.db_test, pytest.mark.need_serialized_dag]
+
+
# when venv tests are run in parallel to other test they create new processes
and this might take
# quite some time in shared docker environment and get some contention even
between different containers
# therefore we have to extend timeouts for those tests
@pytest.mark.execution_timeout(120)
@pytest.mark.virtualenv_operator
-class TestPythonVirtualenvOperator(BaseTestPythonVirtualenvOperator):
+class TestPythonVirtualenvOperator(_VenvTestBase):
+ pytestmark = _VENV_DB_MARKS
opcls = PythonVirtualenvOperator
@staticmethod
@@ -1976,7 +2098,8 @@ class
TestPythonVirtualenvOperator(BaseTestPythonVirtualenvOperator):
# therefore we have to extend timeouts for those tests
@pytest.mark.execution_timeout(120)
@pytest.mark.external_python_operator
-class TestExternalPythonOperator(BaseTestPythonVirtualenvOperator):
+class TestExternalPythonOperator(_VenvTestBase):
+ pytestmark = _VENV_DB_MARKS
opcls = ExternalPythonOperator
@staticmethod
@@ -2338,6 +2461,7 @@ class
BaseTestBranchPythonVirtualenvOperator(BaseTestPythonVirtualenvOperator):
@pytest.mark.execution_timeout(120)
@pytest.mark.virtualenv_operator
class
TestBranchPythonVirtualenvOperator(BaseTestBranchPythonVirtualenvOperator):
+ pytestmark = [pytest.mark.db_test, pytest.mark.need_serialized_dag]
opcls = BranchPythonVirtualenvOperator
@staticmethod
@@ -2356,6 +2480,7 @@ class
TestBranchPythonVirtualenvOperator(BaseTestBranchPythonVirtualenvOperator)
# therefore we have to extend timeouts for those tests
@pytest.mark.external_python_operator
class TestBranchExternalPythonOperator(BaseTestBranchPythonVirtualenvOperator):
+ pytestmark = [pytest.mark.db_test, pytest.mark.need_serialized_dag]
opcls = BranchExternalPythonOperator
@staticmethod
@@ -2368,6 +2493,8 @@ class
TestBranchExternalPythonOperator(BaseTestBranchPythonVirtualenvOperator):
class TestCurrentContext:
+ pytestmark = [pytest.mark.db_test, pytest.mark.need_serialized_dag]
+
def test_current_context_no_context_raise(self):
if AIRFLOW_V_3_0_PLUS:
with pytest.warns(AirflowProviderDeprecationWarning):
@@ -2457,6 +2584,8 @@ def clear_db():
@pytest.mark.usefixtures("clear_db")
class TestCurrentContextRuntime:
+ pytestmark = [pytest.mark.db_test, pytest.mark.need_serialized_dag]
+
def test_context_in_task(self, dag_maker):
with dag_maker(dag_id="assert_context_dag", serialized=True):
op = MyContextAssertOperator(task_id="assert_context")
@@ -2480,6 +2609,8 @@ class TestCurrentContextRuntime:
@pytest.mark.need_serialized_dag(False)
class TestShortCircuitWithTeardown:
+ pytestmark = [pytest.mark.db_test] # keep the class's own
need_serialized_dag(False)
+
@pytest.mark.parametrize(
("ignore_downstream_trigger_rules", "with_teardown", "should_skip",
"expected"),
[
@@ -2611,6 +2742,8 @@ class TestShortCircuitWithTeardown:
class TestPythonAsyncOperator(TestPythonOperator):
+ pytestmark = [pytest.mark.db_test, pytest.mark.need_serialized_dag]
+
def test_run_async_task(self, caplog):
caplog.set_level(logging.INFO, logger=LOGGER_NAME)
diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
index d4617efbbae..9f1c71d4164 100644
--- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
+++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
@@ -2019,6 +2019,7 @@ class InProcessTestSupervisor(ActivitySubprocess):
what: TaskInstance,
task,
logger: FilteringBoundLogger | None = None,
+ client: Client | None = None,
**kwargs,
) -> TaskRunResult:
"""
@@ -2030,6 +2031,9 @@ class InProcessTestSupervisor(ActivitySubprocess):
Dag is already parsed in memory.
Supervisor state and communications are simulated in-memory via
`InProcessSupervisorComms`.
+
+ :param client: optional Execution-API client to use. Defaults to the
DB-backed
+ in-process API server; pass an in-memory/dry-run client to run
without a metadata DB.
"""
# Create supervisor instance
supervisor = cls(
@@ -2037,7 +2041,7 @@ class InProcessTestSupervisor(ActivitySubprocess):
pid=os.getpid(), # Use current process
process=psutil.Process(), # Current process
process_log=logger or
structlog.get_logger(logger_name="task").bind(),
- client=cls._api_client(task.dag),
+ client=client if client is not None else cls._api_client(task.dag),
**kwargs,
)
@@ -2180,10 +2184,15 @@ def set_supervisor_comms(temp_comms):
task_runner.SUPERVISOR_COMMS = old
-def run_task_in_process(ti: TaskInstance, task) -> TaskRunResult:
- """Run a task in-process for testing."""
+def run_task_in_process(ti: TaskInstance, task, client: Client | None = None)
-> TaskRunResult:
+ """
+ Run a task in-process for testing.
+
+ :param client: optional Execution-API client (e.g. an in-memory/dry-run
client to run
+ without a metadata DB). Defaults to the DB-backed in-process API
server.
+ """
# Run the task
- return InProcessTestSupervisor.start(what=ti, task=task)
+ return InProcessTestSupervisor.start(what=ti, task=task, client=client)
# Sockets, even the `.makefile()` function don't correctly do line buffering
on reading. If a chunk is read
