@@ -542,6 +542,80 @@ def test_get_columns_with_constraints(
542542 for either in self .bq_source .yield_table ((table [0 ], table [1 ]))
543543 ]
544544
545+ def test_query_view_names_and_types_with_incremental (self ):
546+ """
547+ Test query_view_names_and_types with incremental extraction enabled
548+ """
549+ self .bq_source .context .get ().__dict__ ["database" ] = MOCK_DB_NAME
550+ self .bq_source .context .get ().__dict__ [
551+ "database_schema"
552+ ] = MOCK_DATABASE_SCHEMA .name .root
553+
554+ mock_view_results = [("view1" ,), ("view2" ,), ("view3" ,)]
555+ mock_materialized_view_results = [("mat_view1" ,), ("mat_view2" ,)]
556+
557+ def mock_execute (query ):
558+ if "INFORMATION_SCHEMA.VIEWS" in query :
559+ return mock_view_results
560+ elif "INFORMATION_SCHEMA.MATERIALIZED_VIEWS" in query :
561+ return mock_materialized_view_results
562+ return []
563+
564+ self .bq_source .engine .execute = mock_execute
565+
566+ mock_processor = Mock ()
567+ mock_processor .get_not_deleted .return_value = ["view1" , "mat_view1" , "view3" ]
568+ self .bq_source .incremental_table_processor = mock_processor
569+ self .bq_source .incremental .enabled = True
570+
571+ result = list (
572+ self .bq_source .query_view_names_and_types (MOCK_DATABASE_SCHEMA .name .root )
573+ )
574+
575+ assert len (result ) == 3
576+ assert result [0 ].name == "view1"
577+ assert result [0 ].type_ == TableType .View
578+ assert result [1 ].name == "view3"
579+ assert result [1 ].type_ == TableType .View
580+ assert result [2 ].name == "mat_view1"
581+ assert result [2 ].type_ == TableType .MaterializedView
582+
583+ mock_processor .get_not_deleted .assert_called ()
584+
585+ def test_query_view_names_and_types_without_incremental (self ):
586+ """
587+ Test query_view_names_and_types without incremental extraction
588+ """
589+ self .bq_source .context .get ().__dict__ ["database" ] = MOCK_DB_NAME
590+ self .bq_source .context .get ().__dict__ [
591+ "database_schema"
592+ ] = MOCK_DATABASE_SCHEMA .name .root
593+
594+ mock_view_results = [("view1" ,), ("view2" ,)]
595+ mock_materialized_view_results = [("mat_view1" ,)]
596+
597+ def mock_execute (query ):
598+ if "INFORMATION_SCHEMA.VIEWS" in query :
599+ return mock_view_results
600+ elif "INFORMATION_SCHEMA.MATERIALIZED_VIEWS" in query :
601+ return mock_materialized_view_results
602+ return []
603+
604+ self .bq_source .engine .execute = mock_execute
605+ self .bq_source .incremental .enabled = False
606+
607+ result = list (
608+ self .bq_source .query_view_names_and_types (MOCK_DATABASE_SCHEMA .name .root )
609+ )
610+
611+ assert len (result ) == 3
612+ assert result [0 ].name == "view1"
613+ assert result [0 ].type_ == TableType .View
614+ assert result [1 ].name == "view2"
615+ assert result [1 ].type_ == TableType .View
616+ assert result [2 ].name == "mat_view1"
617+ assert result [2 ].type_ == TableType .MaterializedView
618+
545619
546620class BigqueryLineageSourceTest (TestCase ):
547621 """
0 commit comments