11import os
2- import time
32import uuid
43from datetime import datetime
54
1615from feast .feature import Feature
1716from feast .feature_table import FeatureTable
1817from feast .value_type import ValueType
18+ from feast .wait import wait_retry_backoff
1919
2020DIR_PATH = os .path .dirname (os .path .realpath (__file__ ))
2121PROJECT_NAME = "basic_" + uuid .uuid4 ().hex .upper ()[0 :6 ]
@@ -236,6 +236,7 @@ def test_get_list_alltypes(
236236 assert actual_list_feature_table == alltypes_featuretable
237237
238238
239+ @pytest .mark .bq
239240def test_ingest (
240241 client : Client ,
241242 customer_entity : Entity ,
@@ -255,12 +256,31 @@ def test_ingest(
255256 client .apply_feature_table (bq_featuretable )
256257 client .ingest (bq_featuretable , bq_dataset , timeout = 120 )
257258
258- # Give time to allow data to propagate to BQ table
259- time .sleep (15 )
260-
259+ from google .api_core .exceptions import NotFound
261260 from google .cloud import bigquery
262261
263262 bq_client = bigquery .Client (project = gcp_project )
263+
264+ # Poll BQ for table until the table has been created
265+ def try_get_table ():
266+ table_exist = False
267+ table_resp = None
268+ try :
269+ table_resp = bq_client .get_table (bq_table_id )
270+
271+ if table_resp and table_resp .table_id == bq_table_id .split ("." )[- 1 ]:
272+ table_exist = True
273+ except NotFound :
274+ pass
275+
276+ return table_resp , table_exist
277+
278+ wait_retry_backoff (
279+ retry_fn = try_get_table ,
280+ timeout_secs = 30 ,
281+ timeout_msg = "Timed out trying to get bigquery table" ,
282+ )
283+
264284 query_string = f"SELECT * FROM `{ bq_table_id } `"
265285
266286 job = bq_client .query (query_string )
0 commit comments