2222from google import auth
2323from vertexai .resources .preview .feature_store import (
2424 offline_store ,
25- FeatureGroup ,
26- Feature ,
2725)
2826
2927pytestmark = [
@@ -266,7 +264,7 @@ def mock_session(bigframes_import_mock):
266264
267265@pytest .fixture ()
268266def mock_fg ():
269- with mock .patch .object (FeatureGroup , "__new__ " ) as mock_fg :
267+ with mock .patch .object (offline_store , "FeatureGroup " ) as mock_fg :
270268 yield mock_fg
271269
272270
@@ -282,12 +280,6 @@ def create_mock_fg(
282280 return fg
283281
284282
285- @pytest .fixture ()
286- def mock_feature ():
287- with mock .patch .object (Feature , "__new__" ) as mock_feature :
288- yield mock_feature
289-
290-
291283def create_mock_feature (
292284 name : str ,
293285 version_column_name : str ,
@@ -360,7 +352,6 @@ def test_one_feature_same_and_different_bq_col_name(
360352 mock_convert_to_bigquery_dataframe ,
361353 mock_session ,
362354 mock_fg ,
363- mock_feature ,
364355 bigframes_import_mock ,
365356):
366357 bigframes , _ = bigframes_import_mock
@@ -382,6 +373,7 @@ def test_one_feature_same_and_different_bq_col_name(
382373 mock_fg .return_value = create_mock_fg (
383374 name = "fake" , entity_id_cols = ["customer_id" ], bq_uri = "bq://my.table"
384375 )
376+ mock_feature = mock .MagicMock ()
385377 mock_feature .return_value = create_mock_feature (
386378 name = feature_name , version_column_name = bq_column_name
387379 )
@@ -407,7 +399,6 @@ def test_one_feature_with_explicit_project_and_location(
407399 mock_convert_to_bigquery_dataframe ,
408400 mock_session ,
409401 mock_fg ,
410- mock_feature ,
411402 bigframes_import_mock ,
412403):
413404 bigframes , _ = bigframes_import_mock
@@ -428,6 +419,7 @@ def test_one_feature_with_explicit_project_and_location(
428419 mock_fg .return_value = create_mock_fg (
429420 name = "fake" , entity_id_cols = ["customer_id" ], bq_uri = "bq://my.table"
430421 )
422+ mock_feature = mock .MagicMock ()
431423 mock_feature .return_value = create_mock_feature (
432424 name = "my_feature" , version_column_name = "my_feature"
433425 )
@@ -457,7 +449,6 @@ def test_one_feature_with_explicit_credentials(
457449 mock_convert_to_bigquery_dataframe ,
458450 mock_session ,
459451 mock_fg ,
460- mock_feature ,
461452 bigframes_import_mock ,
462453):
463454 bigframes , _ = bigframes_import_mock
@@ -478,6 +469,7 @@ def test_one_feature_with_explicit_credentials(
478469 mock_fg .return_value = create_mock_fg (
479470 name = "fake" , entity_id_cols = ["customer_id" ], bq_uri = "bq://my.table"
480471 )
472+ mock_feature = mock .MagicMock ()
481473 mock_feature .return_value = create_mock_feature (
482474 name = "my_feature" , version_column_name = "my_feature"
483475 )
@@ -507,7 +499,6 @@ def test_one_feature_with_explicit_credentials(
507499 # Ensure when getting the FeatureGroup and Feature, the credentials are
508500 # passed through.
509501 mock_fg .assert_called_once_with (
510- FeatureGroup ,
511502 "fake" ,
512503 project = None ,
513504 credentials = credentials ,
0 commit comments