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)