docs(samples): added create feature and create entity type samples an…

…d tests (#984) * feat: SDK feature store samples (create/delete fs) * feat: adding to conftest.py * docs(samples): fixed testing * docs(samples): fixed testing * docs(samples): fixed testing * docs(samples) added changes * docs(samples): style issues * adding create entity * docs(samples): added create feature and entity type * docs(samples): edited test * docs(samples) edited style * moving constants * fixed dates * Update samples/model-builder/create_featurestore_sample_test.py Co-authored-by: Morgan Du <[email protected]> * Update samples/model-builder/test_constants.py Co-authored-by: Morgan Du <[email protected]> * Update samples/model-builder/create_featurestore_sample_test.py Co-authored-by: Morgan Du <[email protected]> * docs(samples): add samples to create/delete featurestore (#980) * feat: SDK feature store samples (create/delete fs) * feat: adding to conftest.py * docs(samples): fixed testing * docs(samples): fixed testing * docs(samples): fixed testing * docs(samples) added changes * docs(samples): style issues * Update samples/model-builder/create_featurestore_sample_test.py Co-authored-by: Morgan Du <[email protected]> * Update samples/model-builder/test_constants.py Co-authored-by: Morgan Du <[email protected]> * Update samples/model-builder/create_featurestore_sample_test.py Co-authored-by: Morgan Du <[email protected]> Co-authored-by: Morgan Du <[email protected]> * Update samples/model-builder/test_constants.py Co-authored-by: Morgan Du <[email protected]> * moving constants * added variables, made fixes, fixed spelling Co-authored-by: Morgan Du <[email protected]>
googleapis · Feb 25, 2022 · d221e6b · d221e6b
1 parent 5fe59a4
commit d221e6b
Show file tree

Hide file tree

Showing 8 changed files with 193 additions and 6 deletions.
diff --git a/samples/model-builder/conftest.py b/samples/model-builder/conftest.py
@@ -379,6 +379,18 @@ def mock_featurestore():
     yield mock
 
 
+@pytest.fixture
+def mock_entity_type():
+    mock = MagicMock(aiplatform.featurestore.EntityType)
+    yield mock
+
+
+@pytest.fixture
+def mock_feature():
+    mock = MagicMock(aiplatform.featurestore.Feature)
+    yield mock
+
+
 @pytest.fixture
 def mock_get_featurestore(mock_featurestore):
     with patch.object(aiplatform.featurestore, "Featurestore") as mock_get_featurestore:
@@ -395,6 +407,24 @@ def mock_create_featurestore(mock_featurestore):
         yield mock_create_featurestore
 
 
+@pytest.fixture
+def mock_create_entity_type(mock_entity_type):
+    with patch.object(
+        aiplatform.featurestore.EntityType, "create"
+    ) as mock_create_entity_type:
+        mock_create_entity_type.return_value = mock_entity_type
+        yield mock_create_entity_type
+
+
+@pytest.fixture
+def mock_create_feature(mock_feature):
+    with patch.object(
+        aiplatform.featurestore.Feature, "create"
+    ) as mock_create_feature:
+        mock_create_feature.return_value = mock_feature
+        yield mock_create_feature
+
+
 @pytest.fixture
 def mock_delete_featurestore(mock_featurestore):
     with patch.object(mock_featurestore, "delete") as mock_delete_featurestore:

diff --git a/samples/model-builder/create_entity_type_sample.py b/samples/model-builder/create_entity_type_sample.py
@@ -0,0 +1,35 @@
+# Copyright 2022 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_entity_type_sample]
+from google.cloud import aiplatform
+
+
+def create_entity_type_sample(
+    project: str, location: str, entity_type_id: str, featurestore_name: str,
+):
+
+    aiplatform.init(project=project, location=location)
+
+    my_entity_type = aiplatform.EntityType.create(
+        entity_type_id=entity_type_id, featurestore_name=featurestore_name
+    )
+
+    my_entity_type.wait()
+
+    return my_entity_type
+
+
+#  [END aiplatform_sdk_create_entity_type_sample]
diff --git a/samples/model-builder/create_entity_type_sample_test.py b/samples/model-builder/create_entity_type_sample_test.py
@@ -0,0 +1,35 @@
+# Copyright 2022 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.
+
+import create_entity_type_sample
+import test_constants as constants
+
+
+def test_create_entity_type_sample(mock_sdk_init, mock_create_entity_type):
+
+    create_entity_type_sample.create_entity_type_sample(
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+        entity_type_id=constants.ENTITY_TYPE_ID,
+        featurestore_name=constants.FEATURESTORE_NAME,
+    )
+
+    mock_sdk_init.assert_called_once_with(
+        project=constants.PROJECT, location=constants.LOCATION
+    )
+
+    mock_create_entity_type.assert_called_once_with(
+        entity_type_id=constants.ENTITY_TYPE_ID,
+        featurestore_name=constants.FEATURESTORE_NAME,
+    )
diff --git a/samples/model-builder/create_feature_sample.py b/samples/model-builder/create_feature_sample.py
@@ -0,0 +1,43 @@
+# Copyright 2022 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_sample]
+from google.cloud import aiplatform
+
+
+def create_feature_sample(
+    project: str,
+    location: str,
+    feature_id: str,
+    value_type: str,
+    entity_type_id: str,
+    featurestore_id: str,
+):
+
+    aiplatform.init(project=project, location=location)
+
+    my_feature = aiplatform.Feature.create(
+        feature_id=feature_id,
+        value_type=value_type,
+        entity_type_name=entity_type_id,
+        featurestore_id=featurestore_id,
+    )
+
+    my_feature.wait()
+
+    return my_feature
+
+
+#  [END aiplatform_sdk_create_feature_sample]
diff --git a/samples/model-builder/create_feature_sample_test.py b/samples/model-builder/create_feature_sample_test.py
@@ -0,0 +1,39 @@
+# Copyright 2022 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.
+
+import create_feature_sample
+import test_constants as constants
+
+
+def test_create_feature_sample(mock_sdk_init, mock_create_feature):
+
+    create_feature_sample.create_feature_sample(
+        project=constants.PROJECT,
+        location=constants.LOCATION,
+        feature_id=constants.FEATURE_ID,
+        value_type=constants.FEATURE_VALUE_TYPE,
+        entity_type_id=constants.ENTITY_TYPE_ID,
+        featurestore_id=constants.FEATURESTORE_ID,
+    )
+
+    mock_sdk_init.assert_called_once_with(
+        project=constants.PROJECT, location=constants.LOCATION
+    )
+
+    mock_create_feature.assert_called_once_with(
+        feature_id=constants.FEATURE_ID,
+        value_type=constants.FEATURE_VALUE_TYPE,
+        entity_type_name=constants.ENTITY_TYPE_ID,
+        featurestore_id=constants.FEATURESTORE_ID,
+    )
diff --git a/samples/model-builder/create_featurestore_sample_test.py b/samples/model-builder/create_featurestore_sample_test.py
@@ -21,7 +21,7 @@ def test_create_featurestore_sample(mock_sdk_init, mock_create_featurestore):
     create_featurestore_sample.create_featurestore_sample(
         project=constants.PROJECT,
         location=constants.LOCATION,
-        featurestore_id=constants.FEAUTURESTORE_ID,
+        featurestore_id=constants.FEATURESTORE_ID,
         online_store_fixed_node_count=constants.ONLINE_STORE_FIXED_NODE_COUNT,
         sync=constants.SYNC,
     )
@@ -31,7 +31,7 @@ def test_create_featurestore_sample(mock_sdk_init, mock_create_featurestore):
     )
 
     mock_create_featurestore.assert_called_once_with(
-        featurestore_id=constants.FEAUTURESTORE_ID,
+        featurestore_id=constants.FEATURESTORE_ID,
         online_store_fixed_node_count=constants.ONLINE_STORE_FIXED_NODE_COUNT,
         sync=constants.SYNC,
     )
diff --git a/samples/model-builder/delete_featurestore_sample_test.py b/samples/model-builder/delete_featurestore_sample_test.py
@@ -23,7 +23,7 @@ def test_delete_featurestore_sample(
     delete_featurestore_sample.delete_featurestore_sample(
         project=constants.PROJECT,
         location=constants.LOCATION,
-        featurestore_name=constants.FEAUTURESTORE_NAME,
+        featurestore_name=constants.FEATURESTORE_NAME,
         sync=constants.SYNC,
         force=constants.FORCE,
     )
@@ -33,7 +33,7 @@ def test_delete_featurestore_sample(
     )
 
     mock_get_featurestore.assert_called_once_with(
-        featurestore_name=constants.FEAUTURESTORE_NAME
+        featurestore_name=constants.FEATURESTORE_NAME
     )
 
     mock_delete_featurestore.assert_called_once_with(

diff --git a/samples/model-builder/test_constants.py b/samples/model-builder/test_constants.py
@@ -199,8 +199,13 @@
 MODEL_TYPE = "CLOUD"
 
 # Feature store constants
-FEAUTURESTORE_NAME = "projects/123/locations/us-central1/featurestores/featurestore_id"
-FEAUTURESTORE_ID = "featurestore_id"
+FEATURESTORE_ID = "featurestore_id"
+FEATURESTORE_NAME = "projects/123/locations/us-central1/featurestores/featurestore_id"
+ENTITY_TYPE_ID = "entity_type_id"
+ENTITY_TYPE_NAME = "projects/123/locations/us-central1/featurestores/featurestore_id/entityTypes/entity_type_id"
+FEATURE_ID = "feature_id"
+FEATURE_NAME = "projects/123/locations/us-central1/featurestores/featurestore_id/entityTypes/entity_type_id/features/feature_id"
+FEATURE_VALUE_TYPE = "INT64"
 ONLINE_STORE_FIXED_NODE_COUNT = 1
 SYNC = True
 FORCE = True