[GitHub] [airflow] kaxil commented on a change in pull request #17577: Have the dag_maker fixture (optionally) give SerializedDAGs

[GitBox <gi...@apache.org>]

kaxil commented on a change in pull request #17577:
URL: https://github.com/apache/airflow/pull/17577#discussion_r687693111



##########
File path: tests/conftest.py
##########
@@ -451,34 +451,60 @@ def dag_maker(request):
 
     The dag_maker.create_dagrun takes the same arguments as dag.create_dagrun
 
+    If you want to operate on serialized DAGs, then either pass ``serialized=True` to the ``dag_maker()``
+    call, or you can mark your test/class/file with ``@pytest.mark.need_serialized_dag(True)``. In both of
+    these cases the ``dag`` returned by the context manager will be a lazily-evaluated proxy object to the
+    SerializedDAG.
     """
-    from airflow.models import DAG, DagModel
+    import lazy_object_proxy
+
+    from airflow.models import DAG, DagBag, DagModel, DagRun, TaskInstance
+    from airflow.models.serialized_dag import SerializedDagModel
     from airflow.utils import timezone
     from airflow.utils.session import provide_session
     from airflow.utils.state import State
 
     DEFAULT_DATE = timezone.datetime(2016, 1, 1)
 
+    want_serialized = False
+    #
+    # Allow changing default serialized behaviour with `@ptest.mark.need_serialized_dag(True)`
+    serialized_marker = request.node.get_closest_marker("need_serialized_dag")
+    if serialized_marker:
+        want_serialized = serialized_marker.args[0]
+
     class DagFactory:
+        def __init__(self):
+            # Keep all the serialized dags we've created in this test
+            self.dagbag = DagBag(os.devnull, include_examples=False, read_dags_from_db=False)
+
         def __enter__(self):
             self.dag.__enter__()
+            if self.want_serialized:
+                return lazy_object_proxy.Proxy(self._serialized_dag)
             return self.dag
 
+        def _serialized_dag(self):
+            return self.serilized_model.dag

Review comment:
       Wait.. looks like a typo in serialized




-- 
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: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org