Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
copybara-service merged 1 commit into from
May 5, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 7 additions & 7 deletions tests/unit/agentplatform/test_feature.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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:
Expand Down
32 changes: 16 additions & 16 deletions tests/unit/agentplatform/test_feature_group.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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(
Expand All @@ -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(
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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}
Expand Down
14 changes: 7 additions & 7 deletions tests/unit/agentplatform/test_feature_monitor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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():
Expand All @@ -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,
)

Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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()

Expand Down
26 changes: 13 additions & 13 deletions tests/unit/agentplatform/test_feature_online_store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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})
Expand Down
Loading
Loading