[Python] airflow with pytest to test xcom data makes an error

I made CustomOperator, which is SampleOperator and the Operator made xcom_push. So, I want to make pytest code testing the xcom data is successfully saved xcom data.

Airflow version: 2.10.5


Here is my full code:

from airflow.models.baseoperator import BaseOperator
import datetime

import pendulum
import pytest

from airflow import DAG
from airflow.utils.state import DagRunState, TaskInstanceState
from airflow.utils.types import DagRunType

DATA_INTERVAL_START = pendulum.now("UTC")
DATA_INTERVAL_END = DATA_INTERVAL_START + datetime.timedelta(days=1)

TEST_DAG_ID = "my_custom_operator_dag"
TEST_TASK_ID = "my_custom_operator_task"

class SampleOperator(BaseOperator):
template_fields = ("_start_date", "_end_date")

def __init__(self, start_date, end_date, **kwargs):
super().__init__(**kwargs)
self._start_date = start_date
self._end_date = end_date

def execute(self, context):
context["ti"].xcom_push(key="start_date", value=self._start_date)
context["ti"].xcom_push(key="end_date", value=self._end_date)
return context


@pytest.fixture()
def dag():
with DAG(
dag_id=TEST_DAG_ID,
schedule="@daily",
start_date=DATA_INTERVAL_START,
) as dag:
SampleOperator(
start_date="{{ prev_ds }}",
end_date="{{ ds }}",
task_id=TEST_TASK_ID
)
return dag


def test_my_custom_operator_execute_no_trigger(dag):
dagrun = dag.create_dagrun(
state=DagRunState.RUNNING,
execution_date=DATA_INTERVAL_START,
data_interval=(DATA_INTERVAL_START, DATA_INTERVAL_END),
start_date=DATA_INTERVAL_END,
run_type=DagRunType.MANUAL,
)
ti = dagrun.get_task_instance(
task_id=TEST_TASK_ID
)

ti.task = dag.get_task(task_id=TEST_TASK_ID)
ti.run(ignore_ti_state=True)
assert ti.state == TaskInstanceState.SUCCESS


But I got an error Here.

TraceBack:

Traceback (most recent call last):
File "/home/airflow/.local/lib/python3.12/site-packages/airflow/utils/json.py", line 92, in default
return serialize(o)
^^^^^^^^^^^^
File "/home/airflow/.local/lib/python3.12/site-packages/airflow/serialization/serde.py", line 190, in serialize
raise TypeError(f"cannot serialize object of type {cls}")
TypeError: cannot serialize object of type <class 'airflow.utils.context.Context'>

....

File "/usr/local/lib/python3.12/json/encoder.py", line 180, in default
raise TypeError(f'Object of type {o.__class__.__name__} '
TypeError: Object of type Context is not JSON serializable


I think the problem is

class SampleOperator(BaseOperator):
...

def execute(self, context):
# here is the problem!
print(context)
context["ti"].xcom_push(key="start_date", value=self._start_date)
context["ti"].xcom_push(key="end_date", value=self._end_date)
return context


print(context) gives below:

...

'ti': <TaskInstance: my_custom_operator_dag.my_custom_operator_task manual__2025-05-06T14:38:13.142804+00:00 [running]>

...


how to solve it? thanks.

Continue reading...