From 2fbf5a5704612d3f87e157e9ebeb8f6583386a06 Mon Sep 17 00:00:00 2001
From: A Vertex SDK engineer <vertex-sdk-bot@google.com>
Date: Thu, 20 Jun 2024 19:32:58 -0700
Subject: [PATCH] feat: sample code for Vertex AI Feature Store

PiperOrigin-RevId: 645230368
---
 samples/model-builder/conftest.py             | 15 ++++++++
 .../create_feature_group_sample.py            | 33 +++++++++++++++++
 .../create_feature_group_sample_test.py       | 35 +++++++++++++++++++
 samples/model-builder/test_constants.py       |  1 +
 4 files changed, 84 insertions(+)
 create mode 100644 samples/model-builder/feature_store/create_feature_group_sample.py
 create mode 100644 samples/model-builder/feature_store/create_feature_group_sample_test.py

diff --git a/samples/model-builder/conftest.py b/samples/model-builder/conftest.py
index 16d9057046..55facdb363 100644
--- a/samples/model-builder/conftest.py
+++ b/samples/model-builder/conftest.py
@@ -716,6 +716,21 @@ def mock_create_optimized_public_online_store(mock_feature_online_store):
         yield mock_create_optimized_store
 
 
+@pytest.fixture
+def mock_feature_group():
+    mock = MagicMock(preview_resources.FeatureGroup)
+    yield mock
+
+
+@pytest.fixture
+def mock_create_feature_group(mock_feature_group):
+    with patch.object(
+        preview_resources.FeatureGroup, "create"
+    ) as mock_create_feature_group:
+        mock_create_feature_group.return_value = mock_feature_group
+        yield mock_create_feature_group
+
+
 @pytest.fixture
 def mock_create_optimized_private_online_store(mock_feature_online_store):
     with patch.object(
diff --git a/samples/model-builder/feature_store/create_feature_group_sample.py b/samples/model-builder/feature_store/create_feature_group_sample.py
new file mode 100644
index 0000000000..7365b60cff
--- /dev/null
+++ b/samples/model-builder/feature_store/create_feature_group_sample.py
@@ -0,0 +1,33 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+#  [START aiplatform_sdk_create_feature_group_sample]
+
+from google.cloud import aiplatform
+from vertexai.resources.preview import feature_store
+
+
+def create_feature_group_sample(
+    project: str,
+    location: str,
+    feature_group_id: str,
+):
+    aiplatform.init(project=project, location=location)
+    fg = feature_store.FeatureGroup.create(
+        feature_group_id
+    )
+    return fg
+
+
+#  [END aiplatform_sdk_create_feature_group_sample]
diff --git a/samples/model-builder/feature_store/create_feature_group_sample_test.py b/samples/model-builder/feature_store/create_feature_group_sample_test.py
new file mode 100644
index 0000000000..9beb6c0816
--- /dev/null
+++ b/samples/model-builder/feature_store/create_feature_group_sample_test.py
@@ -0,0 +1,35 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from feature_store import create_feature_group_sample
+
+import test_constants as constants
+
+
+def test_create_feature_group_sample(
+    mock_sdk_init, mock_create_feature_group
+):
+    create_feature_group_sample.create_feature_group_sample(
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+        feature_group_id=constants.FEATURE_GROUP_ID,
+    )
+
+    mock_sdk_init.assert_called_once_with(
+        project=constants.PROJECT, location=constants.LOCATION
+    )
+
+    mock_create_feature_group.assert_called_once_with(
+        constants.FEATURE_GROUP_ID
+    )
diff --git a/samples/model-builder/test_constants.py b/samples/model-builder/test_constants.py
index 3571e8d9cc..e17baa43a1 100644
--- a/samples/model-builder/test_constants.py
+++ b/samples/model-builder/test_constants.py
@@ -254,6 +254,7 @@
 
 # Feature online store constants
 FEATURE_ONLINE_STORE_ID = "sample_feature_online_store"
+FEATURE_GROUP_ID = "sample_feature_group"
 PROJECT_ALLOWLISTED = ["test-project"]
 
 TABULAR_TARGET_COLUMN = "target_column"