diff --git a/tests/unit/agentplatform/test_feature.py b/tests/unit/agentplatform/test_feature.py index 80882278a4..d967f0dedf 100644 --- a/tests/unit/agentplatform/test_feature.py +++ b/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: diff --git a/tests/unit/agentplatform/test_feature_group.py b/tests/unit/agentplatform/test_feature_group.py index 795d71f950..633f6de1f2 100644 --- a/tests/unit/agentplatform/test_feature_group.py +++ b/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} diff --git a/tests/unit/agentplatform/test_feature_monitor.py b/tests/unit/agentplatform/test_feature_monitor.py index 9779cb5c36..217d8d3d02 100644 --- a/tests/unit/agentplatform/test_feature_monitor.py +++ b/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() diff --git a/tests/unit/agentplatform/test_feature_online_store.py b/tests/unit/agentplatform/test_feature_online_store.py index 12f548c8db..cc6672bd76 100644 --- a/tests/unit/agentplatform/test_feature_online_store.py +++ b/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}) diff --git a/tests/unit/agentplatform/test_feature_view.py b/tests/unit/agentplatform/test_feature_view.py index aa3c5da2a4..6eda8b82b1 100644 --- a/tests/unit/agentplatform/test_feature_view.py +++ b/tests/unit/agentplatform/test_feature_view.py @@ -218,13 +218,13 @@ def test_init_with_fv_id_and_no_fos_id_raises_error(get_fv_mock): + " feature_online_store_id." ), ): - FeatureView(_TEST_FV1_ID) + FeatureView(name=_TEST_FV1_ID) def test_init_with_fv_id(get_fv_mock): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) - fv = FeatureView(_TEST_FV1_ID, feature_online_store_id=_TEST_BIGTABLE_FOS1_ID) + fv = FeatureView(name=_TEST_FV1_ID, feature_online_store_id=_TEST_BIGTABLE_FOS1_ID) get_fv_mock.assert_called_once_with( name=_TEST_FV1_PATH, @@ -244,7 +244,7 @@ def test_init_with_fv_id(get_fv_mock): def test_init_with_fv_path(get_fv_mock): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) - fv = FeatureView(_TEST_FV1_PATH) + fv = FeatureView(name=_TEST_FV1_PATH) get_fv_mock.assert_called_once_with( name=_TEST_FV1_PATH, @@ -332,7 +332,7 @@ def test_delete(delete_fv_mock, fv_logger_mock, get_fos_mock, get_fv_mock, sync= def test_get_sync(get_fv_mock, get_fv_sync_mock): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) - fv_sync = FeatureView(_TEST_FV1_PATH).get_sync(_TEST_FV_SYNC1_ID) + fv_sync = FeatureView(name=_TEST_FV1_PATH).get_sync(_TEST_FV_SYNC1_ID) get_fv_mock.assert_called_once_with( name=_TEST_FV1_PATH, @@ -356,7 +356,7 @@ def test_get_sync(get_fv_mock, get_fv_sync_mock): def test_list_syncs(get_fv_mock, list_fv_syncs_mock): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) - fv_syncs = FeatureView(_TEST_FV1_PATH).list_syncs() + fv_syncs = FeatureView(name=_TEST_FV1_PATH).list_syncs() get_fv_mock.assert_called_once_with( name=_TEST_FV1_PATH, @@ -386,7 +386,7 @@ def test_list_syncs(get_fv_mock, list_fv_syncs_mock): def test_on_demand_sync(get_fv_mock, get_fv_sync_mock, sync_fv_sync_mock): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) - fv_sync = FeatureView(_TEST_FV1_PATH).sync() + fv_sync = FeatureView(name=_TEST_FV1_PATH).sync() get_fv_mock.assert_called_once_with( name=_TEST_FV1_PATH, @@ -416,12 +416,12 @@ def test_fetch_feature_values_bigtable( get_fos_mock, get_fv_mock, fetch_feature_values_mock, fv_logger_mock, output_type ): if output_type == "dict": - fv_dict = FeatureView(_TEST_FV1_PATH).read(key=["key1"]).to_dict() + fv_dict = FeatureView(name=_TEST_FV1_PATH).read(key=["key1"]).to_dict() assert fv_dict == { "features": [{"name": "key1", "value": {"string_value": "value1"}}] } elif output_type == "proto": - fv_proto = FeatureView(_TEST_FV1_PATH).read(key=["key1"]).to_proto() + fv_proto = FeatureView(name=_TEST_FV1_PATH).read(key=["key1"]).to_proto() assert fv_proto == _TEST_FV_FETCH1 fv_logger_mock.assert_has_calls( @@ -440,12 +440,16 @@ def test_fetch_feature_values_optimized( output_type, ): if output_type == "dict": - fv_dict = FeatureView(_TEST_OPTIMIZED_FV1_PATH).read(key=["key1"]).to_dict() + fv_dict = ( + FeatureView(name=_TEST_OPTIMIZED_FV1_PATH).read(key=["key1"]).to_dict() + ) assert fv_dict == { "features": [{"name": "key1", "value": {"string_value": "value1"}}] } elif output_type == "proto": - fv_proto = FeatureView(_TEST_OPTIMIZED_FV1_PATH).read(key=["key1"]).to_proto() + fv_proto = ( + FeatureView(name=_TEST_OPTIMIZED_FV1_PATH).read(key=["key1"]).to_proto() + ) assert fv_proto == _TEST_FV_FETCH1 fv_logger_mock.assert_has_calls( @@ -471,7 +475,7 @@ def test_fetch_feature_values_optimized_no_endpoint( "to complete." ), ): - FeatureView(_TEST_OPTIMIZED_FV2_PATH).read(key=["key1"]).to_dict() + FeatureView(name=_TEST_OPTIMIZED_FV2_PATH).read(key=["key1"]).to_dict() def test_ffv_optimized_psc_with_no_connection_options_raises_error( @@ -479,7 +483,7 @@ def test_ffv_optimized_psc_with_no_connection_options_raises_error( get_optimized_fv_mock, ): with pytest.raises(ValueError) as excinfo: - FeatureView(_TEST_OPTIMIZED_FV1_PATH).read(key=["key1"]) + FeatureView(name=_TEST_OPTIMIZED_FV1_PATH).read(key=["key1"]) assert str(excinfo.value) == ( "Use `connection_options` to specify an IP address. Required for optimized online store with private service connect." @@ -491,7 +495,7 @@ def test_ffv_optimized_psc_with_no_connection_transport_raises_error( get_optimized_fv_mock, ): with pytest.raises(ValueError) as excinfo: - FeatureView(_TEST_OPTIMIZED_FV1_PATH).read( + FeatureView(name=_TEST_OPTIMIZED_FV1_PATH).read( key=["key1"], connection_options=fs_utils.ConnectionOptions( host="1.2.3.4", transport=None @@ -508,7 +512,7 @@ def test_ffv_optimized_psc_with_bad_connection_transport_raises_error( get_optimized_fv_mock, ): with pytest.raises(ValueError) as excinfo: - FeatureView(_TEST_OPTIMIZED_FV1_PATH).read( + FeatureView(name=_TEST_OPTIMIZED_FV1_PATH).read( key=["key1"], connection_options=fs_utils.ConnectionOptions( host="1.2.3.4", transport="hi" @@ -529,7 +533,7 @@ def test_ffv_optimized_psc( fetch_feature_values_mock, output_type, ): - rsp = FeatureView(_TEST_OPTIMIZED_FV1_PATH).read( + rsp = FeatureView(name=_TEST_OPTIMIZED_FV1_PATH).read( key=["key1"], connection_options=fs_utils.ConnectionOptions( host="1.2.3.4", @@ -694,7 +698,7 @@ def test_ffv_optimized_psc_reuse_client_for_same_connection_options_in_same_ffv( grpc_insecure_channel_mock, fetch_feature_values_mock, ): - fv = FeatureView(_TEST_OPTIMIZED_FV1_PATH) + fv = FeatureView(name=_TEST_OPTIMIZED_FV1_PATH) fv.read( key=["key1"], connection_options=fs_utils.ConnectionOptions( @@ -732,7 +736,7 @@ def test_ffv_optimized_psc_different_client_for_different_connection_options( grpc_chan2 = mock.MagicMock(spec=grpc.Channel) grpc_insecure_channel_mock.side_effect = [grpc_chan1, grpc_chan2] - fv = FeatureView(_TEST_OPTIMIZED_FV1_PATH) + fv = FeatureView(name=_TEST_OPTIMIZED_FV1_PATH) fv.read( key=["key1"], connection_options=fs_utils.ConnectionOptions( @@ -763,7 +767,7 @@ def test_ffv_optimized_psc_bad_gapic_client_raises_error( get_psc_optimized_fos_mock, get_optimized_fv_mock, utils_client_with_override_mock ): with pytest.raises(ValueError) as excinfo: - FeatureView(_TEST_OPTIMIZED_FV1_PATH).read( + FeatureView(name=_TEST_OPTIMIZED_FV1_PATH).read( key=["key1"], connection_options=fs_utils.ConnectionOptions( host="1.1.1.1", @@ -787,7 +791,7 @@ def test_search_nearest_entities( if output_type == "dict": fv_dict = ( # Test with entity_id input. - FeatureView(_TEST_EMBEDDING_FV1_PATH) + FeatureView(name=_TEST_EMBEDDING_FV1_PATH) .search( entity_id="key1", neighbor_count=2, @@ -805,7 +809,7 @@ def test_search_nearest_entities( elif output_type == "proto": fv_proto = ( # Test with embedding_value input. - FeatureView(_TEST_EMBEDDING_FV1_PATH) + FeatureView(name=_TEST_EMBEDDING_FV1_PATH) .search(embedding_value=[0.1, 0.2, 0.3]) .to_proto() ) @@ -828,7 +832,7 @@ def test_search_nearest_entities_without_entity_id_or_embedding( fv_logger_mock, ): try: - FeatureView(_TEST_EMBEDDING_FV1_PATH).search().to_proto() + FeatureView(name=_TEST_EMBEDDING_FV1_PATH).search().to_proto() assert not search_nearest_entities_mock.called except ValueError as e: error_msg = ( @@ -844,7 +848,7 @@ def test_search_nearest_entities_no_endpoint( ): """Tests that the public endpoint is not created for the optimized online store.""" try: - FeatureView(_TEST_OPTIMIZED_FV2_PATH).search(entity_id="key1").to_dict() + FeatureView(name=_TEST_OPTIMIZED_FV2_PATH).search(entity_id="key1").to_dict() assert not fetch_feature_values_mock.called except fs_utils.PublicEndpointNotFoundError as e: assert isinstance(e, fs_utils.PublicEndpointNotFoundError) diff --git a/tests/unit/agentplatform/test_model_monitors.py b/tests/unit/agentplatform/test_model_monitors.py index 842f4ae7c3..4f1706a78e 100644 --- a/tests/unit/agentplatform/test_model_monitors.py +++ b/tests/unit/agentplatform/test_model_monitors.py @@ -620,7 +620,7 @@ def test_constructor_creates_client(self, create_client_mock): location=_TEST_LOCATION, credentials=_TEST_CREDENTIALS, ) - ModelMonitor(_TEST_MODEL_MONITOR_ID) + ModelMonitor(model_monitor_name=_TEST_MODEL_MONITOR_ID) create_client_mock.assert_any_call( client_class=utils.ModelMonitoringClientWithOverride, credentials=initializer.global_config.credentials, @@ -634,7 +634,9 @@ def test_constructor_create_client_with_custom_location(self, create_client_mock location=_TEST_LOCATION, credentials=_TEST_CREDENTIALS, ) - ModelMonitor(_TEST_MODEL_MONITOR_ID, location=_TEST_LOCATION_2) + ModelMonitor( + model_monitor_name=_TEST_MODEL_MONITOR_ID, location=_TEST_LOCATION_2 + ) create_client_mock.assert_any_call( client_class=utils.ModelMonitoringClientWithOverride, credentials=initializer.global_config.credentials, @@ -646,7 +648,7 @@ def test_constructor_creates_client_with_custom_credentials( self, create_client_mock ): creds = auth_credentials.AnonymousCredentials() - ModelMonitor(_TEST_MODEL_MONITOR_ID, credentials=creds) + ModelMonitor(model_monitor_name=_TEST_MODEL_MONITOR_ID, credentials=creds) create_client_mock.assert_any_call( client_class=utils.ModelMonitoringClientWithOverride, credentials=creds, diff --git a/tests/unit/vertexai/test_vertexai_feature.py b/tests/unit/vertexai/test_vertexai_feature.py index 151b926438..6da9e3b4fb 100644 --- a/tests/unit/vertexai/test_vertexai_feature.py +++ b/tests/unit/vertexai/test_vertexai_feature.py @@ -120,7 +120,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): @@ -132,13 +132,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, @@ -162,7 +162,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, @@ -185,7 +185,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, @@ -209,7 +209,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, @@ -262,7 +262,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: diff --git a/tests/unit/vertexai/test_vertexai_feature_group.py b/tests/unit/vertexai/test_vertexai_feature_group.py index 5767e1eb9b..c7e1f54b86 100644 --- a/tests/unit/vertexai/test_vertexai_feature_group.py +++ b/tests/unit/vertexai/test_vertexai_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} diff --git a/tests/unit/vertexai/test_vertexai_feature_monitor.py b/tests/unit/vertexai/test_vertexai_feature_monitor.py index 71191319b0..f6934b7e13 100644 --- a/tests/unit/vertexai/test_vertexai_feature_monitor.py +++ b/tests/unit/vertexai/test_vertexai_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() diff --git a/tests/unit/vertexai/test_vertexai_feature_online_store.py b/tests/unit/vertexai/test_vertexai_feature_online_store.py index 5df9e0b339..2964c51296 100644 --- a/tests/unit/vertexai/test_vertexai_feature_online_store.py +++ b/tests/unit/vertexai/test_vertexai_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}) diff --git a/tests/unit/vertexai/test_vertexai_feature_view.py b/tests/unit/vertexai/test_vertexai_feature_view.py index 724959ab95..5fbd17f8ba 100644 --- a/tests/unit/vertexai/test_vertexai_feature_view.py +++ b/tests/unit/vertexai/test_vertexai_feature_view.py @@ -224,7 +224,7 @@ def test_init_with_fv_id_and_no_fos_id_raises_error(get_fv_mock): def test_init_with_fv_id(get_fv_mock): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) - fv = FeatureView(_TEST_FV1_ID, feature_online_store_id=_TEST_BIGTABLE_FOS1_ID) + fv = FeatureView(name=_TEST_FV1_ID, feature_online_store_id=_TEST_BIGTABLE_FOS1_ID) get_fv_mock.assert_called_once_with( name=_TEST_FV1_PATH, @@ -244,7 +244,7 @@ def test_init_with_fv_id(get_fv_mock): def test_init_with_fv_path(get_fv_mock): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) - fv = FeatureView(_TEST_FV1_PATH) + fv = FeatureView(name=_TEST_FV1_PATH) get_fv_mock.assert_called_once_with( name=_TEST_FV1_PATH, @@ -332,7 +332,7 @@ def test_delete(delete_fv_mock, fv_logger_mock, get_fos_mock, get_fv_mock, sync= def test_get_sync(get_fv_mock, get_fv_sync_mock): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) - fv_sync = FeatureView(_TEST_FV1_PATH).get_sync(_TEST_FV_SYNC1_ID) + fv_sync = FeatureView(name=_TEST_FV1_PATH).get_sync(_TEST_FV_SYNC1_ID) get_fv_mock.assert_called_once_with( name=_TEST_FV1_PATH, @@ -356,7 +356,7 @@ def test_get_sync(get_fv_mock, get_fv_sync_mock): def test_list_syncs(get_fv_mock, list_fv_syncs_mock): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) - fv_syncs = FeatureView(_TEST_FV1_PATH).list_syncs() + fv_syncs = FeatureView(name=_TEST_FV1_PATH).list_syncs() get_fv_mock.assert_called_once_with( name=_TEST_FV1_PATH, @@ -386,7 +386,7 @@ def test_list_syncs(get_fv_mock, list_fv_syncs_mock): def test_on_demand_sync(get_fv_mock, get_fv_sync_mock, sync_fv_sync_mock): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) - fv_sync = FeatureView(_TEST_FV1_PATH).sync() + fv_sync = FeatureView(name=_TEST_FV1_PATH).sync() get_fv_mock.assert_called_once_with( name=_TEST_FV1_PATH, @@ -416,12 +416,12 @@ def test_fetch_feature_values_bigtable( get_fos_mock, get_fv_mock, fetch_feature_values_mock, fv_logger_mock, output_type ): if output_type == "dict": - fv_dict = FeatureView(_TEST_FV1_PATH).read(key=["key1"]).to_dict() + fv_dict = FeatureView(name=_TEST_FV1_PATH).read(key=["key1"]).to_dict() assert fv_dict == { "features": [{"name": "key1", "value": {"string_value": "value1"}}] } elif output_type == "proto": - fv_proto = FeatureView(_TEST_FV1_PATH).read(key=["key1"]).to_proto() + fv_proto = FeatureView(name=_TEST_FV1_PATH).read(key=["key1"]).to_proto() assert fv_proto == _TEST_FV_FETCH1 fv_logger_mock.assert_has_calls( @@ -440,12 +440,16 @@ def test_fetch_feature_values_optimized( output_type, ): if output_type == "dict": - fv_dict = FeatureView(_TEST_OPTIMIZED_FV1_PATH).read(key=["key1"]).to_dict() + fv_dict = ( + FeatureView(name=_TEST_OPTIMIZED_FV1_PATH).read(key=["key1"]).to_dict() + ) assert fv_dict == { "features": [{"name": "key1", "value": {"string_value": "value1"}}] } elif output_type == "proto": - fv_proto = FeatureView(_TEST_OPTIMIZED_FV1_PATH).read(key=["key1"]).to_proto() + fv_proto = ( + FeatureView(name=_TEST_OPTIMIZED_FV1_PATH).read(key=["key1"]).to_proto() + ) assert fv_proto == _TEST_FV_FETCH1 fv_logger_mock.assert_has_calls( @@ -471,7 +475,7 @@ def test_fetch_feature_values_optimized_no_endpoint( "to complete." ), ): - FeatureView(_TEST_OPTIMIZED_FV2_PATH).read(key=["key1"]).to_dict() + FeatureView(name=_TEST_OPTIMIZED_FV2_PATH).read(key=["key1"]).to_dict() def test_ffv_optimized_psc_with_no_connection_options_raises_error( @@ -479,7 +483,7 @@ def test_ffv_optimized_psc_with_no_connection_options_raises_error( get_optimized_fv_mock, ): with pytest.raises(ValueError) as excinfo: - FeatureView(_TEST_OPTIMIZED_FV1_PATH).read(key=["key1"]) + FeatureView(name=_TEST_OPTIMIZED_FV1_PATH).read(key=["key1"]) assert str(excinfo.value) == ( "Use `connection_options` to specify an IP address. Required for optimized online store with private service connect." @@ -491,7 +495,7 @@ def test_ffv_optimized_psc_with_no_connection_transport_raises_error( get_optimized_fv_mock, ): with pytest.raises(ValueError) as excinfo: - FeatureView(_TEST_OPTIMIZED_FV1_PATH).read( + FeatureView(name=_TEST_OPTIMIZED_FV1_PATH).read( key=["key1"], connection_options=fs_utils.ConnectionOptions( host="1.2.3.4", transport=None @@ -508,7 +512,7 @@ def test_ffv_optimized_psc_with_bad_connection_transport_raises_error( get_optimized_fv_mock, ): with pytest.raises(ValueError) as excinfo: - FeatureView(_TEST_OPTIMIZED_FV1_PATH).read( + FeatureView(name=_TEST_OPTIMIZED_FV1_PATH).read( key=["key1"], connection_options=fs_utils.ConnectionOptions( host="1.2.3.4", transport="hi" @@ -529,7 +533,7 @@ def test_ffv_optimized_psc( fetch_feature_values_mock, output_type, ): - rsp = FeatureView(_TEST_OPTIMIZED_FV1_PATH).read( + rsp = FeatureView(name=_TEST_OPTIMIZED_FV1_PATH).read( key=["key1"], connection_options=fs_utils.ConnectionOptions( host="1.2.3.4", @@ -694,7 +698,7 @@ def test_ffv_optimized_psc_reuse_client_for_same_connection_options_in_same_ffv( grpc_insecure_channel_mock, fetch_feature_values_mock, ): - fv = FeatureView(_TEST_OPTIMIZED_FV1_PATH) + fv = FeatureView(name=_TEST_OPTIMIZED_FV1_PATH) fv.read( key=["key1"], connection_options=fs_utils.ConnectionOptions( @@ -732,7 +736,7 @@ def test_ffv_optimized_psc_different_client_for_different_connection_options( grpc_chan2 = mock.MagicMock(spec=grpc.Channel) grpc_insecure_channel_mock.side_effect = [grpc_chan1, grpc_chan2] - fv = FeatureView(_TEST_OPTIMIZED_FV1_PATH) + fv = FeatureView(name=_TEST_OPTIMIZED_FV1_PATH) fv.read( key=["key1"], connection_options=fs_utils.ConnectionOptions( @@ -763,7 +767,7 @@ def test_ffv_optimized_psc_bad_gapic_client_raises_error( get_psc_optimized_fos_mock, get_optimized_fv_mock, utils_client_with_override_mock ): with pytest.raises(ValueError) as excinfo: - FeatureView(_TEST_OPTIMIZED_FV1_PATH).read( + FeatureView(name=_TEST_OPTIMIZED_FV1_PATH).read( key=["key1"], connection_options=fs_utils.ConnectionOptions( host="1.1.1.1", @@ -787,7 +791,7 @@ def test_search_nearest_entities( if output_type == "dict": fv_dict = ( # Test with entity_id input. - FeatureView(_TEST_EMBEDDING_FV1_PATH) + FeatureView(name=_TEST_EMBEDDING_FV1_PATH) .search( entity_id="key1", neighbor_count=2, @@ -805,7 +809,7 @@ def test_search_nearest_entities( elif output_type == "proto": fv_proto = ( # Test with embedding_value input. - FeatureView(_TEST_EMBEDDING_FV1_PATH) + FeatureView(name=_TEST_EMBEDDING_FV1_PATH) .search(embedding_value=[0.1, 0.2, 0.3]) .to_proto() ) @@ -828,7 +832,7 @@ def test_search_nearest_entities_without_entity_id_or_embedding( fv_logger_mock, ): try: - FeatureView(_TEST_EMBEDDING_FV1_PATH).search().to_proto() + FeatureView(name=_TEST_EMBEDDING_FV1_PATH).search().to_proto() assert not search_nearest_entities_mock.called except ValueError as e: error_msg = ( @@ -844,7 +848,7 @@ def test_search_nearest_entities_no_endpoint( ): """Tests that the public endpoint is not created for the optimized online store.""" try: - FeatureView(_TEST_OPTIMIZED_FV2_PATH).search(entity_id="key1").to_dict() + FeatureView(name=_TEST_OPTIMIZED_FV2_PATH).search(entity_id="key1").to_dict() assert not fetch_feature_values_mock.called except fs_utils.PublicEndpointNotFoundError as e: assert isinstance(e, fs_utils.PublicEndpointNotFoundError) diff --git a/tests/unit/vertexai/test_vertexai_model_monitors.py b/tests/unit/vertexai/test_vertexai_model_monitors.py index 11a02d0236..9c6138f387 100644 --- a/tests/unit/vertexai/test_vertexai_model_monitors.py +++ b/tests/unit/vertexai/test_vertexai_model_monitors.py @@ -620,7 +620,7 @@ def test_constructor_creates_client(self, create_client_mock): location=_TEST_LOCATION, credentials=_TEST_CREDENTIALS, ) - ModelMonitor(_TEST_MODEL_MONITOR_ID) + ModelMonitor(model_monitor_name=_TEST_MODEL_MONITOR_ID) create_client_mock.assert_any_call( client_class=utils.ModelMonitoringClientWithOverride, credentials=initializer.global_config.credentials, @@ -634,7 +634,9 @@ def test_constructor_create_client_with_custom_location(self, create_client_mock location=_TEST_LOCATION, credentials=_TEST_CREDENTIALS, ) - ModelMonitor(_TEST_MODEL_MONITOR_ID, location=_TEST_LOCATION_2) + ModelMonitor( + model_monitor_name=_TEST_MODEL_MONITOR_ID, location=_TEST_LOCATION_2 + ) create_client_mock.assert_any_call( client_class=utils.ModelMonitoringClientWithOverride, credentials=initializer.global_config.credentials, @@ -646,7 +648,7 @@ def test_constructor_creates_client_with_custom_credentials( self, create_client_mock ): creds = auth_credentials.AnonymousCredentials() - ModelMonitor(_TEST_MODEL_MONITOR_ID, credentials=creds) + ModelMonitor(model_monitor_name=_TEST_MODEL_MONITOR_ID, credentials=creds) create_client_mock.assert_any_call( client_class=utils.ModelMonitoringClientWithOverride, credentials=creds,