diff --git a/samples/model-builder/conftest.py b/samples/model-builder/conftest.py index d18fca064d..64cfc85902 100644 --- 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 new file mode 100644 index 0000000000..eee0ddb8f5 --- /dev/null +++ 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 new file mode 100644 index 0000000000..fe7dfd3f77 --- /dev/null +++ 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 new file mode 100644 index 0000000000..b9b2328ae8 --- /dev/null +++ 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 new file mode 100644 index 0000000000..673f475691 --- /dev/null +++ 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 index e133f4582b..2531f58c6d 100644 --- 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 index 5361c8c881..21b7cba84c 100644 --- 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 index ebd65211c3..50d02ec7a0 100644 --- 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