1818 BigQueryOfflineStoreConfig ,
1919 BigQueryRetrievalJob ,
2020 _bq_create_entity_union_table ,
21-
2221)
23- from feast .infra .offline_stores .offline_utils import gather_all_entities
2422from feast .infra .offline_stores .bigquery_source import BigQuerySource
25- from feast .infra .offline_stores .offline_utils import FeatureViewQueryContext
23+ from feast .infra .offline_stores .offline_utils import (
24+ FeatureViewQueryContext ,
25+ gather_all_entities ,
26+ )
2627from feast .infra .online_stores .sqlite import SqliteOnlineStoreConfig
2728from feast .repo_config import RepoConfig
2829
@@ -145,9 +146,7 @@ def _make_client(self) -> MagicMock:
145146
146147 @patch ("feast.infra.offline_stores.bigquery.block_until_done" )
147148 @patch ("feast.infra.offline_stores.bigquery._utc_now" )
148- def test_single_view_creates_table_with_correct_sql (
149- self , mock_utc_now , mock_block
150- ):
149+ def test_single_view_creates_table_with_correct_sql (self , mock_utc_now , mock_block ):
151150 mock_utc_now .return_value = END
152151 client = self ._make_client ()
153152 fv = _make_feature_view_mock ("fv1" , ["customer_id" ])
@@ -271,7 +270,9 @@ def test_timestamp_filter_uses_start_and_end(self, mock_utc_now, mock_block):
271270
272271 sql = client .query .call_args [0 ][0 ]
273272 assert "2023-01-01T00:00:00" in sql # start_date
274- assert "2024-01-01T00:00:00" in sql # end_date (appears twice: WHERE + entity_ts value)
273+ assert (
274+ "2024-01-01T00:00:00" in sql
275+ ) # end_date (appears twice: WHERE + entity_ts value)
275276 assert "BETWEEN TIMESTAMP(" in sql
276277
277278 @patch ("feast.infra.offline_stores.bigquery.block_until_done" )
@@ -319,9 +320,9 @@ def test_sets_table_expiry(self, mock_utc_now, mock_block):
319320 client .update_table .assert_called_once ()
320321 updated_table = client .update_table .call_args [0 ][0 ]
321322 # expiry should be 30 minutes after _utc_now()
322- assert updated_table .expires == datetime (2024 , 1 , 1 , tzinfo = timezone . utc ) + timedelta (
323- minutes = 30
324- )
323+ assert updated_table .expires == datetime (
324+ 2024 , 1 , 1 , tzinfo = timezone . utc
325+ ) + timedelta ( minutes = 30 )
325326
326327
327328# ---------------------------------------------------------------------------
@@ -338,7 +339,9 @@ def repo_config():
338339def mock_bq_client ():
339340 client = MagicMock ()
340341 client .project = "my-project"
341- client .query .return_value = MagicMock (state = "DONE" , exception = lambda timeout = None : None )
342+ client .query .return_value = MagicMock (
343+ state = "DONE" , exception = lambda timeout = None : None
344+ )
342345 client .get_table .return_value = MagicMock ()
343346 return client
344347
@@ -348,8 +351,12 @@ class TestGetHistoricalFeaturesNonEntityMode:
348351 @patch ("feast.infra.offline_stores.bigquery._utc_now" )
349352 @patch ("feast.infra.offline_stores.bigquery._bq_create_entity_union_table" )
350353 @patch ("feast.infra.offline_stores.bigquery._upload_entity_df" )
351- @patch ("feast.infra.offline_stores.bigquery.offline_utils.get_feature_view_query_context" )
352- @patch ("feast.infra.offline_stores.bigquery.offline_utils.build_point_in_time_query" )
354+ @patch (
355+ "feast.infra.offline_stores.bigquery.offline_utils.get_feature_view_query_context"
356+ )
357+ @patch (
358+ "feast.infra.offline_stores.bigquery.offline_utils.build_point_in_time_query"
359+ )
353360 @patch ("feast.infra.offline_stores.bigquery._get_table_reference_for_new_entity" )
354361 @patch ("feast.infra.offline_stores.bigquery._get_bigquery_client" )
355362 def test_non_entity_mode_calls_union_table_not_upload (
@@ -403,8 +410,12 @@ def test_non_entity_mode_calls_union_table_not_upload(
403410 @patch ("feast.infra.offline_stores.bigquery._utc_now" )
404411 @patch ("feast.infra.offline_stores.bigquery._bq_create_entity_union_table" )
405412 @patch ("feast.infra.offline_stores.bigquery._upload_entity_df" )
406- @patch ("feast.infra.offline_stores.bigquery.offline_utils.get_feature_view_query_context" )
407- @patch ("feast.infra.offline_stores.bigquery.offline_utils.build_point_in_time_query" )
413+ @patch (
414+ "feast.infra.offline_stores.bigquery.offline_utils.get_feature_view_query_context"
415+ )
416+ @patch (
417+ "feast.infra.offline_stores.bigquery.offline_utils.build_point_in_time_query"
418+ )
408419 @patch ("feast.infra.offline_stores.bigquery._get_table_reference_for_new_entity" )
409420 @patch ("feast.infra.offline_stores.bigquery._get_bigquery_client" )
410421 def test_non_entity_mode_uses_entity_ts_as_timestamp_col (
@@ -455,8 +466,12 @@ def test_non_entity_mode_uses_entity_ts_as_timestamp_col(
455466 @patch ("feast.infra.offline_stores.bigquery._utc_now" )
456467 @patch ("feast.infra.offline_stores.bigquery._bq_create_entity_union_table" )
457468 @patch ("feast.infra.offline_stores.bigquery._upload_entity_df" )
458- @patch ("feast.infra.offline_stores.bigquery.offline_utils.get_feature_view_query_context" )
459- @patch ("feast.infra.offline_stores.bigquery.offline_utils.build_point_in_time_query" )
469+ @patch (
470+ "feast.infra.offline_stores.bigquery.offline_utils.get_feature_view_query_context"
471+ )
472+ @patch (
473+ "feast.infra.offline_stores.bigquery.offline_utils.build_point_in_time_query"
474+ )
460475 @patch ("feast.infra.offline_stores.bigquery._get_table_reference_for_new_entity" )
461476 @patch ("feast.infra.offline_stores.bigquery._get_bigquery_client" )
462477 def test_non_entity_mode_passes_start_and_end_to_union_table (
@@ -507,8 +522,12 @@ def test_non_entity_mode_passes_start_and_end_to_union_table(
507522 @patch ("feast.infra.offline_stores.bigquery.block_until_done" )
508523 @patch ("feast.infra.offline_stores.bigquery._utc_now" )
509524 @patch ("feast.infra.offline_stores.bigquery._bq_create_entity_union_table" )
510- @patch ("feast.infra.offline_stores.bigquery.offline_utils.get_feature_view_query_context" )
511- @patch ("feast.infra.offline_stores.bigquery.offline_utils.build_point_in_time_query" )
525+ @patch (
526+ "feast.infra.offline_stores.bigquery.offline_utils.get_feature_view_query_context"
527+ )
528+ @patch (
529+ "feast.infra.offline_stores.bigquery.offline_utils.build_point_in_time_query"
530+ )
512531 @patch ("feast.infra.offline_stores.bigquery._get_table_reference_for_new_entity" )
513532 @patch ("feast.infra.offline_stores.bigquery._get_bigquery_client" )
514533 def test_non_entity_mode_returns_retrieval_job (
@@ -553,8 +572,12 @@ def test_non_entity_mode_returns_retrieval_job(
553572 @patch ("feast.infra.offline_stores.bigquery.block_until_done" )
554573 @patch ("feast.infra.offline_stores.bigquery._utc_now" )
555574 @patch ("feast.infra.offline_stores.bigquery._bq_create_entity_union_table" )
556- @patch ("feast.infra.offline_stores.bigquery.offline_utils.get_feature_view_query_context" )
557- @patch ("feast.infra.offline_stores.bigquery.offline_utils.build_point_in_time_query" )
575+ @patch (
576+ "feast.infra.offline_stores.bigquery.offline_utils.get_feature_view_query_context"
577+ )
578+ @patch (
579+ "feast.infra.offline_stores.bigquery.offline_utils.build_point_in_time_query"
580+ )
558581 @patch ("feast.infra.offline_stores.bigquery._get_table_reference_for_new_entity" )
559582 @patch ("feast.infra.offline_stores.bigquery._get_bigquery_client" )
560583 def test_non_entity_mode_metadata_excludes_timestamp_col_from_keys (
@@ -610,9 +633,15 @@ class TestGetHistoricalFeaturesEntityDfMode:
610633 @patch ("feast.infra.offline_stores.bigquery._get_entity_df_event_timestamp_range" )
611634 @patch ("feast.infra.offline_stores.bigquery._get_entity_schema" )
612635 @patch ("feast.infra.offline_stores.bigquery.offline_utils.get_expected_join_keys" )
613- @patch ("feast.infra.offline_stores.bigquery.offline_utils.assert_expected_columns_in_entity_df" )
614- @patch ("feast.infra.offline_stores.bigquery.offline_utils.get_feature_view_query_context" )
615- @patch ("feast.infra.offline_stores.bigquery.offline_utils.build_point_in_time_query" )
636+ @patch (
637+ "feast.infra.offline_stores.bigquery.offline_utils.assert_expected_columns_in_entity_df"
638+ )
639+ @patch (
640+ "feast.infra.offline_stores.bigquery.offline_utils.get_feature_view_query_context"
641+ )
642+ @patch (
643+ "feast.infra.offline_stores.bigquery.offline_utils.build_point_in_time_query"
644+ )
616645 @patch ("feast.infra.offline_stores.bigquery._get_table_reference_for_new_entity" )
617646 @patch ("feast.infra.offline_stores.bigquery._get_bigquery_client" )
618647 def test_entity_df_mode_calls_upload_not_union_table (
@@ -671,9 +700,15 @@ def test_entity_df_mode_calls_upload_not_union_table(
671700 @patch ("feast.infra.offline_stores.bigquery._get_entity_df_event_timestamp_range" )
672701 @patch ("feast.infra.offline_stores.bigquery._get_entity_schema" )
673702 @patch ("feast.infra.offline_stores.bigquery.offline_utils.get_expected_join_keys" )
674- @patch ("feast.infra.offline_stores.bigquery.offline_utils.assert_expected_columns_in_entity_df" )
675- @patch ("feast.infra.offline_stores.bigquery.offline_utils.get_feature_view_query_context" )
676- @patch ("feast.infra.offline_stores.bigquery.offline_utils.build_point_in_time_query" )
703+ @patch (
704+ "feast.infra.offline_stores.bigquery.offline_utils.assert_expected_columns_in_entity_df"
705+ )
706+ @patch (
707+ "feast.infra.offline_stores.bigquery.offline_utils.get_feature_view_query_context"
708+ )
709+ @patch (
710+ "feast.infra.offline_stores.bigquery.offline_utils.build_point_in_time_query"
711+ )
677712 @patch ("feast.infra.offline_stores.bigquery._get_table_reference_for_new_entity" )
678713 @patch ("feast.infra.offline_stores.bigquery._get_bigquery_client" )
679714 def test_entity_df_sql_string_mode_works (
0 commit comments