chore: add FeatureGroup list tests

PiperOrigin-RevId: 631927820
googleapis · May 8, 2024 · 378c68a · 378c68a
1 parent 7fea754
commit 378c68a
Show file tree

Hide file tree

Showing 2 changed files with 96 additions and 0 deletions.
diff --git a/tests/unit/vertexai/feature_store_constants.py b/tests/unit/vertexai/feature_store_constants.py
@@ -278,3 +278,39 @@
     ),
     labels=_TEST_FG1_LABELS,
 )
+
+
+_TEST_FG2_ID = "my_fg2"
+_TEST_FG2_PATH = f"{_TEST_PARENT}/featureGroups/{_TEST_FG2_ID}"
+_TEST_FG2_BQ_URI = f"bq://{_TEST_PROJECT}.my_dataset.my_table_for_fg2"
+_TEST_FG2_ENTITY_ID_COLUMNS = ["entity_id1", "entity_id2"]
+_TEST_FG2_LABELS = {"my_key2": "my_fg2"}
+_TEST_FG2 = types.feature_group.FeatureGroup(
+    name=_TEST_FG2_PATH,
+    big_query=types.feature_group.FeatureGroup.BigQuery(
+        big_query_source=types.io.BigQuerySource(
+            input_uri=_TEST_FG2_BQ_URI,
+        ),
+        entity_id_columns=_TEST_FG2_ENTITY_ID_COLUMNS,
+    ),
+    labels=_TEST_FG2_LABELS,
+)
+
+
+_TEST_FG3_ID = "my_fg3"
+_TEST_FG3_PATH = f"{_TEST_PARENT}/featureGroups/{_TEST_FG3_ID}"
+_TEST_FG3_BQ_URI = f"bq://{_TEST_PROJECT}.my_dataset.my_table_for_fg3"
+_TEST_FG3_ENTITY_ID_COLUMNS = ["entity_id1", "entity_id2", "entity_id3"]
+_TEST_FG3_LABELS = {"my_key3": "my_fg3"}
+_TEST_FG3 = types.feature_group.FeatureGroup(
+    name=_TEST_FG3_PATH,
+    big_query=types.feature_group.FeatureGroup.BigQuery(
+        big_query_source=types.io.BigQuerySource(
+            input_uri=_TEST_FG3_BQ_URI,
+        ),
+        entity_id_columns=_TEST_FG3_ENTITY_ID_COLUMNS,
+    ),
+    labels=_TEST_FG3_LABELS,
+)
+
+_TEST_FG_LIST = [_TEST_FG1, _TEST_FG2, _TEST_FG3]
diff --git a/tests/unit/vertexai/test_feature_group.py b/tests/unit/vertexai/test_feature_group.py
@@ -49,6 +49,17 @@
     _TEST_FG1_BQ_URI,
     _TEST_FG1_ENTITY_ID_COLUMNS,
     _TEST_FG1_LABELS,
+    _TEST_FG2_ID,
+    _TEST_FG2_PATH,
+    _TEST_FG2_BQ_URI,
+    _TEST_FG2_ENTITY_ID_COLUMNS,
+    _TEST_FG2_LABELS,
+    _TEST_FG3_ID,
+    _TEST_FG3_PATH,
+    _TEST_FG3_BQ_URI,
+    _TEST_FG3_ENTITY_ID_COLUMNS,
+    _TEST_FG3_LABELS,
+    _TEST_FG_LIST,
 )
 
 
@@ -87,6 +98,16 @@ def create_fg_mock():
         yield create_fg_mock
 
 
+@pytest.fixture
+def list_fg_mock():
+    with patch.object(
+        feature_registry_service_client.FeatureRegistryServiceClient,
+        "list_feature_groups",
+    ) as list_fg_mock:
+        list_fg_mock.return_value = _TEST_FG_LIST
+        yield list_fg_mock
+
+
 def fg_eq(
     fg_to_check: FeatureGroup,
     name: str,
@@ -233,3 +254,42 @@ def test_create_fg(
         location=_TEST_LOCATION,
         labels=_TEST_FG1_LABELS,
     )
+
+
+def test_list(list_fg_mock):
+    aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
+
+    feature_groups = FeatureGroup.list()
+
+    list_fg_mock.assert_called_once_with(request={"parent": _TEST_PARENT})
+    assert len(feature_groups) == len(_TEST_FG_LIST)
+    fg_eq(
+        feature_groups[0],
+        name=_TEST_FG1_ID,
+        resource_name=_TEST_FG1_PATH,
+        source_uri=_TEST_FG1_BQ_URI,
+        entity_id_columns=_TEST_FG1_ENTITY_ID_COLUMNS,
+        project=_TEST_PROJECT,
+        location=_TEST_LOCATION,
+        labels=_TEST_FG1_LABELS,
+    )
+    fg_eq(
+        feature_groups[1],
+        name=_TEST_FG2_ID,
+        resource_name=_TEST_FG2_PATH,
+        source_uri=_TEST_FG2_BQ_URI,
+        entity_id_columns=_TEST_FG2_ENTITY_ID_COLUMNS,
+        project=_TEST_PROJECT,
+        location=_TEST_LOCATION,
+        labels=_TEST_FG2_LABELS,
+    )
+    fg_eq(
+        feature_groups[2],
+        name=_TEST_FG3_ID,
+        resource_name=_TEST_FG3_PATH,
+        source_uri=_TEST_FG3_BQ_URI,
+        entity_id_columns=_TEST_FG3_ENTITY_ID_COLUMNS,
+        project=_TEST_PROJECT,
+        location=_TEST_LOCATION,
+        labels=_TEST_FG3_LABELS,
+    )