def test_get_sub_dag_multiple_input_views2(self) -> None:
all_views_dag_walker = BigQueryViewDagWalker(self.x_shaped_dag_views_list)
input_views = [
self.x_shaped_dag_views_list[2],
self.x_shaped_dag_views_list[4],
]
# Get descendants sub-dag
sub_dag = all_views_dag_walker.get_descendants_sub_dag(
input_views,
)
# Only should include this view
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, sub_dag.views)
# Get ancestors sub-dag
sub_dag = all_views_dag_walker.get_ancestors_sub_dag(input_views)
expected_views = [
self.x_shaped_dag_views_list[0],
self.x_shaped_dag_views_list[1],
self.x_shaped_dag_views_list[2],
self.x_shaped_dag_views_list[4],
]
self.assertCountEqual(expected_views, sub_dag.views)
def test_sub_dag_with_cycle(self) -> None:
all_views_dag_walker = BigQueryViewDagWalker(self.diamond_shaped_dag_views_list)
input_views = [
self.diamond_shaped_dag_views_list[1],
self.diamond_shaped_dag_views_list[4],
]
# Get descendants sub-dag
sub_dag = all_views_dag_walker.get_descendants_sub_dag(input_views)
expected_views = [
self.diamond_shaped_dag_views_list[1],
self.diamond_shaped_dag_views_list[2],
self.diamond_shaped_dag_views_list[3],
self.diamond_shaped_dag_views_list[4],
self.diamond_shaped_dag_views_list[5],
]
self.assertCountEqual(expected_views, sub_dag.views)
# Get ancestors sub-dag
sub_dag = all_views_dag_walker.get_ancestors_sub_dag(
input_views,
)
expected_views = [
self.diamond_shaped_dag_views_list[0],
self.diamond_shaped_dag_views_list[1],
self.diamond_shaped_dag_views_list[2],
self.diamond_shaped_dag_views_list[4],
]
self.assertCountEqual(expected_views, sub_dag.views)
def test_get_sub_dag_leaf_node(self) -> None:
all_views_dag_walker = BigQueryViewDagWalker(self.x_shaped_dag_views_list)
# Get descendants sub-dag
sub_dag = all_views_dag_walker.get_descendants_sub_dag(
[self.x_shaped_dag_views_list[4]],
)
# Only should include this view
expected_views = [self.x_shaped_dag_views_list[4]]
self.assertCountEqual(expected_views, sub_dag.views)
# Get ancestors sub-dag
sub_dag = all_views_dag_walker.get_ancestors_sub_dag(
[self.x_shaped_dag_views_list[4]]
)
expected_views = [
v
for v in self.x_shaped_dag_views_list
# This view does not depend on other leaf view "table_4"
if v.view_id != "table_4"
]
self.assertCountEqual(expected_views, sub_dag.views)
def test_get_sub_dag_input_views_not_all_in_input_dag(self) -> None:
all_views_dag_walker = BigQueryViewDagWalker(self.x_shaped_dag_views_list[0:2])
with self.assertRaises(ValueError) as e:
# Add different input views to build source dag
_ = all_views_dag_walker.get_descendants_sub_dag(
self.x_shaped_dag_views_list[1:5],
)
self.assertTrue(
str(e.exception).startswith("Found input views not in source DAG:")
)
def test_get_sub_dag_empty_input_views(self) -> None:
all_views_dag_walker = BigQueryViewDagWalker(self.x_shaped_dag_views_list)
# Get descendants sub-dag
sub_dag = all_views_dag_walker.get_descendants_sub_dag([])
self.assertCountEqual([], sub_dag.views)
# Get ancestors sub-dag
sub_dag = all_views_dag_walker.get_ancestors_sub_dag([])
self.assertCountEqual([], sub_dag.views)
def test_get_sub_dag_single_node_input(self) -> None:
all_views_dag_walker = BigQueryViewDagWalker(self.x_shaped_dag_views_list[0:1])
# Get descendants sub-dag
sub_dag = all_views_dag_walker.get_descendants_sub_dag(
self.x_shaped_dag_views_list[0:1],
)
expected_views = self.x_shaped_dag_views_list[0:1]
self.assertCountEqual(expected_views, sub_dag.views)
# Get ancestors sub-dag
sub_dag = all_views_dag_walker.get_ancestors_sub_dag(
self.x_shaped_dag_views_list[0:1]
)
# Only should include this view
expected_views = self.x_shaped_dag_views_list[0:1]
self.assertCountEqual(expected_views, sub_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)
