fix: async test for ingestion

googleapis · Dec 15, 2021 · a5ab277 · a5ab277
1 parent 14bd73a
commit a5ab277
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/google/cloud/aiplatform/featurestore/entity_type.py b/google/cloud/aiplatform/featurestore/entity_type.py
@@ -898,7 +898,7 @@ def _import_feature_values(
 
         return self
 
-    @base.optional_sync()
+    @base.optional_sync(return_input_arg="self")
     def ingest_from_bq(
         self,
         feature_ids: List[str],
@@ -1000,7 +1000,7 @@ def ingest_from_bq(
             request_metadata=request_metadata,
         )
 
-    @base.optional_sync()
+    @base.optional_sync(return_input_arg="self")
     def ingest_from_gcs(
         self,
         feature_ids: List[str],

diff --git a/tests/unit/aiplatform/test_featurestores.py b/tests/unit/aiplatform/test_featurestores.py
@@ -1061,6 +1061,7 @@ def test_ingest_from_bq(self, import_feature_values_mock, sync):
             feature_time=_TEST_FEATURE_TIME_FIELD,
             bq_source_uri=_TEST_BQ_SOURCE_URI,
             feature_source_fields=_TEST_IMPORTING_FEATURE_SOURCE_FIELDS,
+            sync=sync,
         )
 
         if not sync:
@@ -1091,6 +1092,7 @@ def test_ingest_from_gcs(self, import_feature_values_mock, sync):
             feature_time=_TEST_FEATURE_TIME,
             gcs_source_uris=_TEST_GCS_AVRO_SOURCE_URIS,
             gcs_source_type=_TEST_GCS_SOURCE_TYPE_AVRO,
+            sync=sync,
         )
 
         if not sync: