def test_union_dags_same_view_different_object(self) -> None:
view = BigQueryView(
dataset_id="dataset_1",
view_id="table_1",
description="table_1 description",
should_materialize=True,
materialized_address_override=BigQueryAddress(
dataset_id="other_dataset_1", table_id="other_table_1"
),
view_query_template="SELECT * FROM `{project_id}.source_dataset.source_table`",
)
unioned_dag = BigQueryViewDagWalker.union_dags(
BigQueryViewDagWalker([view]),
BigQueryViewDagWalker(
[
BigQueryView(
dataset_id="dataset_1",
view_id="table_1",
description="table_1 description",
should_materialize=True,
materialized_address_override=BigQueryAddress(
dataset_id="other_dataset_1", table_id="other_table_1"
),
view_query_template="SELECT * FROM `{project_id}.source_dataset.source_table`",
)
]
),
)
self.assertCountEqual([view], unioned_dag.views)
def test_union_dags(self) -> None:
x_shaped_dag_walker = BigQueryViewDagWalker(self.x_shaped_dag_views_list)
# This DAG is a superset of the X-shaped DAG
diamond_shaped_dag_walker = BigQueryViewDagWalker(
self.diamond_shaped_dag_views_list
)
unioned_dag = BigQueryViewDagWalker.union_dags(
x_shaped_dag_walker, diamond_shaped_dag_walker
)
self.assertCountEqual(self.diamond_shaped_dag_views_list, unioned_dag.views)
def test_get_sub_dag_middle_node(self) -> None:
all_views_dag_walker = BigQueryViewDagWalker(self.x_shaped_dag_views_list)
# Get descendants sub-dag
descendants_sub_dag = all_views_dag_walker.get_descendants_sub_dag(
[self.x_shaped_dag_views_list[2]],
)
expected_views = [
self.x_shaped_dag_views_list[2],
self.x_shaped_dag_views_list[3],
self.x_shaped_dag_views_list[4],
]
self.assertCountEqual(expected_views, descendants_sub_dag.views)
# Get ancestors sub-dag
ancestors_sub_dag = all_views_dag_walker.get_ancestors_sub_dag(
[self.x_shaped_dag_views_list[2]]
)
expected_views = [
self.x_shaped_dag_views_list[0],
self.x_shaped_dag_views_list[1],
self.x_shaped_dag_views_list[2],
]
self.assertCountEqual(expected_views, ancestors_sub_dag.views)
# Get both directions sub-dag
both_directions_dag = BigQueryViewDagWalker.union_dags(
descendants_sub_dag, ancestors_sub_dag
)
expected_views = self.x_shaped_dag_views_list
self.assertCountEqual(expected_views, both_directions_dag.views)
def test_union_dags_empty(self) -> None:
unioned_dag = BigQueryViewDagWalker.union_dags(
BigQueryViewDagWalker([]), BigQueryViewDagWalker([])
)
self.assertCountEqual([], unioned_dag.views)
