fix: flaky tests in feature store relates cases due to class instantiation

PiperOrigin-RevId: 910747896

Author: yyyu-google

tests/unit/agentplatform/test_feature.py

@@ -118,7 +118,7 @@ def test_init_with_feature_id_and_no_fg_id_raises_error(get_feature_mock):
             + " feature_group_id."
         ),
     ):
-        Feature(_TEST_FG1_F1_ID)
+        Feature(name=_TEST_FG1_F1_ID)
 
 
 def test_init_with_feature_path_and_fg_id_raises_error(get_feature_mock):
@@ -130,13 +130,13 @@ def test_init_with_feature_path_and_fg_id_raises_error(get_feature_mock):
             "Since feature 'projects/test-project/locations/us-central1/featureGroups/my_fg1/features/my_fg1_f1' is provided as a path, feature_group_id should not be specified."
         ),
     ):
-        Feature(_TEST_FG1_F1_PATH, feature_group_id=_TEST_FG1_ID)
+        Feature(name=_TEST_FG1_F1_PATH, feature_group_id=_TEST_FG1_ID)
 
 
 def test_init_with_feature_id(get_feature_mock):
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
 
-    feature = Feature(_TEST_FG1_F1_ID, feature_group_id=_TEST_FG1_ID)
+    feature = Feature(name=_TEST_FG1_F1_ID, feature_group_id=_TEST_FG1_ID)
 
     get_feature_mock.assert_called_once_with(
         name=_TEST_FG1_F1_PATH,
@@ -160,7 +160,7 @@ def test_init_with_feature_id_for_explicit_version_column(
 ):
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
 
-    feature = Feature(_TEST_FG1_F2_ID, feature_group_id=_TEST_FG1_ID)
+    feature = Feature(name=_TEST_FG1_F2_ID, feature_group_id=_TEST_FG1_ID)
 
     get_feature_with_version_column_mock.assert_called_once_with(
         name=_TEST_FG1_F2_PATH,
@@ -183,7 +183,7 @@ def test_init_with_feature_id_for_explicit_version_column(
 def test_init_with_feature_path(get_feature_mock):
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
 
-    feature = Feature(_TEST_FG1_F1_PATH)
+    feature = Feature(name=_TEST_FG1_F1_PATH)
 
     get_feature_mock.assert_called_once_with(
         name=_TEST_FG1_F1_PATH,
@@ -207,7 +207,7 @@ def test_init_with_feature_path_for_explicit_version_column(
 ):
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
 
-    feature = Feature(_TEST_FG1_F2_PATH)
+    feature = Feature(name=_TEST_FG1_F2_PATH)
 
     get_feature_with_version_column_mock.assert_called_once_with(
         name=_TEST_FG1_F2_PATH,
@@ -260,7 +260,7 @@ def test_delete_feature(
 ):
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
 
-    feature = FeatureGroup(_TEST_FG1_ID).get_feature(_TEST_FG1_F1_ID)
+    feature = FeatureGroup(name=_TEST_FG1_ID).get_feature(_TEST_FG1_F1_ID)
     feature.delete(sync=sync)
 
     if not sync:

tests/unit/agentplatform/test_feature_group.py

@@ -252,7 +252,7 @@ def fg_eq(
 def test_init(feature_group_name, get_fg_mock):
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
 
-    fg = FeatureGroup(feature_group_name)
+    fg = FeatureGroup(name=feature_group_name)
 
     get_fg_mock.assert_called_once_with(
         name=_TEST_FG1_PATH,
@@ -415,7 +415,7 @@ def test_list(list_fg_mock):
 def test_delete(delete_fg_mock, get_fg_mock, fg_logger_mock, force, sync):
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
 
-    fg = FeatureGroup(_TEST_FG1_ID)
+    fg = FeatureGroup(name=_TEST_FG1_ID)
     fg.delete(force=force, sync=sync)
 
     if not sync:
@@ -444,7 +444,7 @@ def test_delete(delete_fg_mock, get_fg_mock, fg_logger_mock, force, sync):
 def test_get_feature(get_fg_mock, get_feature_mock):
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
 
-    fg = FeatureGroup(_TEST_FG1_ID)
+    fg = FeatureGroup(name=_TEST_FG1_ID)
     feature = fg.get_feature(_TEST_FG1_F1_ID)
 
     get_feature_mock.assert_called_once_with(
@@ -469,7 +469,7 @@ def test_get_feature_with_latest_stats_count(
 ):
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
 
-    fg = FeatureGroup(_TEST_FG1_ID)
+    fg = FeatureGroup(name=_TEST_FG1_ID)
     feature = fg.get_feature(_TEST_FG1_F1_ID, latest_stats_count=1)
 
     get_feature_with_stats_and_anomalies_mock.assert_called_once_with(
@@ -503,7 +503,7 @@ def test_get_feature_credentials_set_in_init(mock_base_instantiate_client):
     mock_base_instantiate_client.return_value.get_feature_group.return_value = _TEST_FG1
     mock_base_instantiate_client.return_value.get_feature.return_value = _TEST_FG1_F1
 
-    fg = FeatureGroup(_TEST_FG1_ID)
+    fg = FeatureGroup(name=_TEST_FG1_ID)
     mock_base_instantiate_client.assert_called_with(
         location=_TEST_LOCATION,
         credentials=credentials,
@@ -538,7 +538,7 @@ def test_get_feature_from_feature_group_with_explicit_credentials(
     mock_base_instantiate_client.return_value.get_feature.return_value = _TEST_FG1_F1
 
     credentials = mock.MagicMock(spec=auth_credentials.Credentials)
-    fg = FeatureGroup(_TEST_FG1_ID, credentials=credentials)
+    fg = FeatureGroup(name=_TEST_FG1_ID, credentials=credentials)
     mock_base_instantiate_client.assert_called_with(
         location=_TEST_LOCATION,
         credentials=credentials,
@@ -576,7 +576,7 @@ def test_get_feature_from_feature_group_with_explicit_credentials_overrides_init
     mock_base_instantiate_client.return_value.get_feature.return_value = _TEST_FG1_F1
 
     credentials = mock.MagicMock(spec=auth_credentials.Credentials)
-    fg = FeatureGroup(_TEST_FG1_ID, credentials=credentials)
+    fg = FeatureGroup(name=_TEST_FG1_ID, credentials=credentials)
     mock_base_instantiate_client.assert_called_with(
         location=_TEST_LOCATION,
         credentials=credentials,
@@ -608,7 +608,7 @@ def test_get_feature_with_explicit_credentials(mock_base_instantiate_client):
     mock_base_instantiate_client.return_value.get_feature_group.return_value = _TEST_FG1
     mock_base_instantiate_client.return_value.get_feature.return_value = _TEST_FG1_F1
 
-    fg = FeatureGroup(_TEST_FG1_ID)
+    fg = FeatureGroup(name=_TEST_FG1_ID)
     mock_base_instantiate_client.assert_called_with(
         location=_TEST_LOCATION,
         credentials=mock.ANY,
@@ -646,7 +646,7 @@ def test_get_feature_with_explicit_credentials_overrides_init_credentials(
     mock_base_instantiate_client.return_value.get_feature_group.return_value = _TEST_FG1
     mock_base_instantiate_client.return_value.get_feature.return_value = _TEST_FG1_F1
 
-    fg = FeatureGroup(_TEST_FG1_ID)
+    fg = FeatureGroup(name=_TEST_FG1_ID)
     mock_base_instantiate_client.assert_called_with(
         location=_TEST_LOCATION,
         credentials=init_credentials,
@@ -682,7 +682,7 @@ def test_get_feature_with_explicit_credentials_overrides_feature_group_credentia
     mock_base_instantiate_client.return_value.get_feature.return_value = _TEST_FG1_F1
 
     feature_group_credentials = mock.MagicMock(spec=auth_credentials.Credentials)
-    fg = FeatureGroup(_TEST_FG1_ID, credentials=feature_group_credentials)
+    fg = FeatureGroup(name=_TEST_FG1_ID, credentials=feature_group_credentials)
     mock_base_instantiate_client.assert_called_with(
         location=_TEST_LOCATION,
         credentials=feature_group_credentials,
@@ -721,7 +721,7 @@ def test_get_feature_with_explicit_credentials_overrides_init_and_feature_group_
     mock_base_instantiate_client.return_value.get_feature.return_value = _TEST_FG1_F1
 
     feature_group_credentials = mock.MagicMock(spec=auth_credentials.Credentials)
-    fg = FeatureGroup(_TEST_FG1_ID, credentials=feature_group_credentials)
+    fg = FeatureGroup(name=_TEST_FG1_ID, credentials=feature_group_credentials)
     mock_base_instantiate_client.assert_called_with(
         location=_TEST_LOCATION,
         credentials=feature_group_credentials,
@@ -760,7 +760,7 @@ def test_create_feature(
 ):
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
 
-    fg = FeatureGroup(_TEST_FG1_ID)
+    fg = FeatureGroup(name=_TEST_FG1_ID)
     feature = fg.create_feature(
         _TEST_FG1_F1_ID,
         description=_TEST_FG1_F1_DESCRIPTION,
@@ -826,7 +826,7 @@ def test_create_feature_with_version_feature_column(
 ):
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
 
-    fg = FeatureGroup(_TEST_FG1_ID)
+    fg = FeatureGroup(name=_TEST_FG1_ID)
     feature = fg.create_feature(
         _TEST_FG1_F2_ID,
         version_column_name=_TEST_FG1_F2_VERSION_COLUMN_NAME,
@@ -886,7 +886,7 @@ def test_create_feature_with_version_feature_column(
 def test_list_features(get_fg_mock, list_features_mock):
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
 
-    features = FeatureGroup(_TEST_FG1_ID).list_features()
+    features = FeatureGroup(name=_TEST_FG1_ID).list_features()
 
     list_features_mock.assert_called_once_with(request={"parent": _TEST_FG1_PATH})
     assert len(features) == len(_TEST_FG1_FEATURE_LIST)
@@ -923,7 +923,7 @@ def test_create_feature_monitor(
 ):
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
 
-    fg = FeatureGroup(_TEST_FG1_ID)
+    fg = FeatureGroup(name=_TEST_FG1_ID)
     feature_monitor = fg.create_feature_monitor(
         _TEST_FG1_FM1_ID,
         description=_TEST_FG1_FM1_DESCRIPTION,
@@ -997,7 +997,7 @@ def test_list_feature_monitors(
 ):
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
 
-    feature_monitors = FeatureGroup(_TEST_FG1_ID).list_feature_monitors()
+    feature_monitors = FeatureGroup(name=_TEST_FG1_ID).list_feature_monitors()
 
     list_feature_monitors_mock.assert_called_once_with(
         request={"parent": _TEST_FG1_PATH}

tests/unit/agentplatform/test_feature_monitor.py

@@ -158,7 +158,7 @@ def test_init_with_feature_monitor_id_and_no_fg_id_raises_error():
             " specify feature_group_id."
         ),
     ):
-        FeatureMonitor(_TEST_FG1_FM1_ID)
+        FeatureMonitor(name=_TEST_FG1_FM1_ID)
 
 
 def test_init_with_feature_monitor_path_and_fg_id_raises_error():
@@ -176,7 +176,7 @@ def test_init_with_feature_monitor_path_and_fg_id_raises_error():
         ),
     ):
         FeatureMonitor(
-            _TEST_FG1_FM1_PATH,
+            name=_TEST_FG1_FM1_PATH,
             feature_group_id=_TEST_FG1_ID,
         )
 
@@ -216,7 +216,7 @@ def test_init_with_feature_monitor_path(get_feature_monitor_mock):
         location=_TEST_LOCATION,
     )
 
-    feature_monitor = FeatureMonitor(_TEST_FG1_FM1_PATH)
+    feature_monitor = FeatureMonitor(name=_TEST_FG1_FM1_PATH)
 
     get_feature_monitor_mock.assert_called_once_with(
         name=_TEST_FG1_FM1_PATH,
@@ -242,7 +242,7 @@ def test_init_with_feature_monitor_job_path(get_feature_monitor_job_mock):
         location=_TEST_LOCATION,
     )
 
-    feature_monitor_job = FeatureMonitor.FeatureMonitorJob(_TEST_FG1_FMJ1_PATH)
+    feature_monitor_job = FeatureMonitor.FeatureMonitorJob(name=_TEST_FG1_FMJ1_PATH)
 
     get_feature_monitor_job_mock.assert_called_once_with(
         name=_TEST_FG1_FMJ1_PATH,
@@ -274,7 +274,7 @@ def test_create_feature_monitor_job(
     )
 
     fm = FeatureMonitor(
-        _TEST_FG1_FM1_ID,
+        name=_TEST_FG1_FM1_ID,
         feature_group_id=_TEST_FG1_ID,
     )
     feature_monitor_job = fm.create_feature_monitor_job(
@@ -314,7 +314,7 @@ def test_get_feature_monitor_job(
     )
 
     fm = FeatureMonitor(
-        _TEST_FG1_FM1_ID,
+        name=_TEST_FG1_FM1_ID,
         feature_group_id=_TEST_FG1_ID,
     )
     feature_monitor_job = fm.get_feature_monitor_job(_TEST_FG1_FMJ1_ID)
@@ -344,7 +344,7 @@ def test_list_feature_monitors_jobs(
     )
 
     feature_monitor_jobs = FeatureMonitor(
-        _TEST_FG1_FM1_ID,
+        name=_TEST_FG1_FM1_ID,
         feature_group_id=_TEST_FG1_ID,
     ).list_feature_monitor_jobs()
 

tests/unit/agentplatform/test_feature_online_store.py

@@ -150,7 +150,7 @@ def fos_eq(
 def test_init(online_store_name, get_fos_mock):
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
 
-    fos = FeatureOnlineStore(online_store_name)
+    fos = FeatureOnlineStore(name=online_store_name)
 
     get_fos_mock.assert_called_once_with(
         name=_TEST_BIGTABLE_FOS1_PATH, retry=base._DEFAULT_RETRY
@@ -439,7 +439,7 @@ def test_list(list_fos_mock):
 def test_delete(force, delete_fos_mock, get_fos_mock, fos_logger_mock, sync=True):
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
 
-    fos = FeatureOnlineStore(_TEST_BIGTABLE_FOS1_ID)
+    fos = FeatureOnlineStore(name=_TEST_BIGTABLE_FOS1_ID)
     fos.delete(force=force, sync=sync)
 
     if not sync:
@@ -467,7 +467,7 @@ def test_delete(force, delete_fos_mock, get_fos_mock, fos_logger_mock, sync=True
 
 def test_create_fv_none_source_raises_error(get_fos_mock):
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
-    fos = FeatureOnlineStore(_TEST_BIGTABLE_FOS1_ID)
+    fos = FeatureOnlineStore(name=_TEST_BIGTABLE_FOS1_ID)
 
     with pytest.raises(
         ValueError,
@@ -478,7 +478,7 @@ def test_create_fv_none_source_raises_error(get_fos_mock):
 
 def test_create_fv_wrong_object_type_raises_error(get_fos_mock):
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
-    fos = FeatureOnlineStore(_TEST_BIGTABLE_FOS1_ID)
+    fos = FeatureOnlineStore(name=_TEST_BIGTABLE_FOS1_ID)
 
     with pytest.raises(
         ValueError,
@@ -492,7 +492,7 @@ def test_create_fv_wrong_object_type_raises_error(get_fos_mock):
 
 def test_create_bq_fv_bad_uri_raises_error(get_fos_mock):
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
-    fos = FeatureOnlineStore(_TEST_BIGTABLE_FOS1_ID)
+    fos = FeatureOnlineStore(name=_TEST_BIGTABLE_FOS1_ID)
 
     with pytest.raises(
         ValueError,
@@ -509,7 +509,7 @@ def test_create_bq_fv_bad_entity_id_columns_raises_error(
     entity_id_columns, get_fos_mock
 ):
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
-    fos = FeatureOnlineStore(_TEST_BIGTABLE_FOS1_ID)
+    fos = FeatureOnlineStore(name=_TEST_BIGTABLE_FOS1_ID)
 
     with pytest.raises(
         ValueError,
@@ -526,7 +526,7 @@ def test_create_bq_fv_bad_entity_id_columns_raises_error(
 )
 def test_create_fr_fv_invalid_feature_name_raises_error(features, get_fos_mock):
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
-    fos = FeatureOnlineStore(_TEST_BIGTABLE_FOS1_ID)
+    fos = FeatureOnlineStore(name=_TEST_BIGTABLE_FOS1_ID)
 
     with pytest.raises(
         ValueError,
@@ -551,7 +551,7 @@ def test_create_bq_fv(
     fos_logger_mock,
 ):
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
-    fos = FeatureOnlineStore(_TEST_BIGTABLE_FOS1_ID)
+    fos = FeatureOnlineStore(name=_TEST_BIGTABLE_FOS1_ID)
 
     fv = fos.create_feature_view(
         _TEST_FV1_ID,
@@ -613,7 +613,7 @@ def test_create_embedding_fv(
     get_optimized_embedding_fv_mock,
 ):
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
-    fos = FeatureOnlineStore(_TEST_ESF_OPTIMIZED_FOS_ID)
+    fos = FeatureOnlineStore(name=_TEST_ESF_OPTIMIZED_FOS_ID)
 
     embedding_fv = fos.create_feature_view(
         _TEST_OPTIMIZED_EMBEDDING_FV_ID,
@@ -639,7 +639,7 @@ def test_create_embedding_fv(
 
 def test_create_rag_fv_bad_uri_raises_error(get_fos_mock):
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
-    fos = FeatureOnlineStore(_TEST_BIGTABLE_FOS1_ID)
+    fos = FeatureOnlineStore(name=_TEST_BIGTABLE_FOS1_ID)
 
     with pytest.raises(
         ValueError,
@@ -662,7 +662,7 @@ def test_create_rag_fv(
     fos_logger_mock,
 ):
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
-    fos = FeatureOnlineStore(_TEST_BIGTABLE_FOS1_ID)
+    fos = FeatureOnlineStore(name=_TEST_BIGTABLE_FOS1_ID)
 
     rag_fv = fos.create_feature_view(
         _TEST_FV3_ID,
@@ -729,7 +729,7 @@ def test_create_registry_fv(
     fos_logger_mock,
 ):
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
-    fos = FeatureOnlineStore(_TEST_BIGTABLE_FOS1_ID)
+    fos = FeatureOnlineStore(name=_TEST_BIGTABLE_FOS1_ID)
 
     registry_fv = fos.create_feature_view(
         _TEST_FV4_ID,
@@ -807,7 +807,7 @@ def test_list_feature_views(
 ):
     aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
 
-    fos = FeatureOnlineStore(_TEST_BIGTABLE_FOS1_ID)
+    fos = FeatureOnlineStore(name=_TEST_BIGTABLE_FOS1_ID)
     feature_views = fos.list_feature_views()
 
     list_fv_mock.assert_called_once_with(request={"parent": _TEST_BIGTABLE_FOS1_PATH})