ashb commented on a change in pull request #20931:
URL: https://github.com/apache/airflow/pull/20931#discussion_r787673672
##########
File path: airflow/serialization/enums.py
##########
@@ -48,3 +48,4 @@ class DagAttributeTypes(str, Enum):
TASK_GROUP = 'taskgroup'
EDGE_INFO = 'edgeinfo'
PARAM = 'param'
+ XCOM_REF = 'xcom_ref'
Review comment:
Good point -- better to be consistent.
##########
File path: tests/serialization/test_dag_serialization.py
##########
@@ -1590,14 +1593,58 @@ def test_mapped_operator_serde():
{"__type": "dict", "__var": {'a': 'b'}},
]
},
- 'partial_kwargs': {},
+ 'partial_kwargs': {
+ 'executor_config': {
+ '__type': 'dict',
+ '__var': {
+ 'dict': {"__type": "dict", "__var": {'sub': 'value'}},
+ },
+ },
+ },
'task_id': 'a',
'template_fields': ['bash_command', 'env'],
}
op = SerializedBaseOperator.deserialize_operator(serialized)
+ assert isinstance(op, MappedOperator)
assert op.operator_class == "airflow.operators.bash.BashOperator"
+ assert op.mapped_kwargs['bash_command'] == literal
+ assert op.partial_kwargs['executor_config'] == {'dict': {'sub': 'value'}}
+
+
+def test_mapped_operator_xcomarg_serde():
+ from airflow.models.xcom_arg import XComArg
+
+ with DAG("test-dag", start_date=datetime(2020, 1, 1)) as dag:
+ task1 = BaseOperator(task_id="op1")
+ xcomarg = XComArg(task1, "test_key")
+ mapped = MockOperator(task_id='task_2').map(arg2=xcomarg)
+
+ serialized = SerializedBaseOperator._serialize(mapped)
+ assert serialized == {
+ '_is_dummy': False,
+ '_is_mapped': True,
+ '_task_module': 'tests.test_utils.mock_operators',
+ '_task_type': 'MockOperator',
+ 'downstream_task_ids': [],
+ 'mapped_kwargs': {'arg2': {'__type': 'xcom_ref', '__var': {'task_id':
'op1', 'key': 'test_key'}}},
Review comment:
```suggestion
'mapped_kwargs': {'arg2': {'__type': 'xcomref', '__var': {'task_id':
'op1', 'key': 'test_key'}}},
```
--
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]
