Skip to content

Commit d221e6b

Browse files
nayaknishantmorgandu
nayaknishant
and
morgandu
authored
docs(samples): added create feature and create entity type samples and 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 <morgandu@google.com> * Update samples/model-builder/test_constants.py Co-authored-by: Morgan Du <morgandu@google.com> * Update samples/model-builder/create_featurestore_sample_test.py Co-authored-by: Morgan Du <morgandu@google.com> * 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 <morgandu@google.com> * Update samples/model-builder/test_constants.py Co-authored-by: Morgan Du <morgandu@google.com> * Update samples/model-builder/create_featurestore_sample_test.py Co-authored-by: Morgan Du <morgandu@google.com> Co-authored-by: Morgan Du <morgandu@google.com> * Update samples/model-builder/test_constants.py Co-authored-by: Morgan Du <morgandu@google.com> * moving constants * added variables, made fixes, fixed spelling Co-authored-by: Morgan Du <morgandu@google.com>
1 parent 5fe59a4 commit d221e6b

8 files changed

Lines changed: 193 additions & 6 deletions

‎samples/model-builder/conftest.py‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -379,6 +379,18 @@ def mock_featurestore():
379379
yield mock
380380

381381

382+
@pytest.fixture
383+
def mock_entity_type():
384+
mock = MagicMock(aiplatform.featurestore.EntityType)
385+
yield mock
386+
387+
388+
@pytest.fixture
389+
def mock_feature():
390+
mock = MagicMock(aiplatform.featurestore.Feature)
391+
yield mock
392+
393+
382394
@pytest.fixture
383395
def mock_get_featurestore(mock_featurestore):
384396
with patch.object(aiplatform.featurestore, "Featurestore") as mock_get_featurestore:
@@ -395,6 +407,24 @@ def mock_create_featurestore(mock_featurestore):
395407
yield mock_create_featurestore
396408

397409

410+
@pytest.fixture
411+
def mock_create_entity_type(mock_entity_type):
412+
with patch.object(
413+
aiplatform.featurestore.EntityType, "create"
414+
) as mock_create_entity_type:
415+
mock_create_entity_type.return_value = mock_entity_type
416+
yield mock_create_entity_type
417+
418+
419+
@pytest.fixture
420+
def mock_create_feature(mock_feature):
421+
with patch.object(
422+
aiplatform.featurestore.Feature, "create"
423+
) as mock_create_feature:
424+
mock_create_feature.return_value = mock_feature
425+
yield mock_create_feature
426+
427+
398428
@pytest.fixture
399429
def mock_delete_featurestore(mock_featurestore):
400430
with patch.object(mock_featurestore, "delete") as mock_delete_featurestore:

‎samples/model-builder/create_entity_type_sample.py‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
# Copyright 2022 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
16+
# [START aiplatform_sdk_create_entity_type_sample]
17+
from google.cloud import aiplatform
18+
19+
20+
def create_entity_type_sample(
21+
project: str, location: str, entity_type_id: str, featurestore_name: str,
22+
):
23+
24+
aiplatform.init(project=project, location=location)
25+
26+
my_entity_type = aiplatform.EntityType.create(
27+
entity_type_id=entity_type_id, featurestore_name=featurestore_name
28+
)
29+
30+
my_entity_type.wait()
31+
32+
return my_entity_type
33+
34+
35+
# [END aiplatform_sdk_create_entity_type_sample]

‎samples/model-builder/create_entity_type_sample_test.py‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
# Copyright 2022 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import create_entity_type_sample
16+
import test_constants as constants
17+
18+
19+
def test_create_entity_type_sample(mock_sdk_init, mock_create_entity_type):
20+
21+
create_entity_type_sample.create_entity_type_sample(
22+
project=constants.PROJECT,
23+
location=constants.LOCATION,
24+
entity_type_id=constants.ENTITY_TYPE_ID,
25+
featurestore_name=constants.FEATURESTORE_NAME,
26+
)
27+
28+
mock_sdk_init.assert_called_once_with(
29+
project=constants.PROJECT, location=constants.LOCATION
30+
)
31+
32+
mock_create_entity_type.assert_called_once_with(
33+
entity_type_id=constants.ENTITY_TYPE_ID,
34+
featurestore_name=constants.FEATURESTORE_NAME,
35+
)

‎samples/model-builder/create_feature_sample.py‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
# Copyright 2022 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
16+
# [START aiplatform_sdk_create_feature_sample]
17+
from google.cloud import aiplatform
18+
19+
20+
def create_feature_sample(
21+
project: str,
22+
location: str,
23+
feature_id: str,
24+
value_type: str,
25+
entity_type_id: str,
26+
featurestore_id: str,
27+
):
28+
29+
aiplatform.init(project=project, location=location)
30+
31+
my_feature = aiplatform.Feature.create(
32+
feature_id=feature_id,
33+
value_type=value_type,
34+
entity_type_name=entity_type_id,
35+
featurestore_id=featurestore_id,
36+
)
37+
38+
my_feature.wait()
39+
40+
return my_feature
41+
42+
43+
# [END aiplatform_sdk_create_feature_sample]

‎samples/model-builder/create_feature_sample_test.py‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
# Copyright 2022 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import create_feature_sample
16+
import test_constants as constants
17+
18+
19+
def test_create_feature_sample(mock_sdk_init, mock_create_feature):
20+
21+
create_feature_sample.create_feature_sample(
22+
project=constants.PROJECT,
23+
location=constants.LOCATION,
24+
feature_id=constants.FEATURE_ID,
25+
value_type=constants.FEATURE_VALUE_TYPE,
26+
entity_type_id=constants.ENTITY_TYPE_ID,
27+
featurestore_id=constants.FEATURESTORE_ID,
28+
)
29+
30+
mock_sdk_init.assert_called_once_with(
31+
project=constants.PROJECT, location=constants.LOCATION
32+
)
33+
34+
mock_create_feature.assert_called_once_with(
35+
feature_id=constants.FEATURE_ID,
36+
value_type=constants.FEATURE_VALUE_TYPE,
37+
entity_type_name=constants.ENTITY_TYPE_ID,
38+
featurestore_id=constants.FEATURESTORE_ID,
39+
)

‎samples/model-builder/create_featurestore_sample_test.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ def test_create_featurestore_sample(mock_sdk_init, mock_create_featurestore):
2121
create_featurestore_sample.create_featurestore_sample(
2222
project=constants.PROJECT,
2323
location=constants.LOCATION,
24-
featurestore_id=constants.FEAUTURESTORE_ID,
24+
featurestore_id=constants.FEATURESTORE_ID,
2525
online_store_fixed_node_count=constants.ONLINE_STORE_FIXED_NODE_COUNT,
2626
sync=constants.SYNC,
2727
)
@@ -31,7 +31,7 @@ def test_create_featurestore_sample(mock_sdk_init, mock_create_featurestore):
3131
)
3232

3333
mock_create_featurestore.assert_called_once_with(
34-
featurestore_id=constants.FEAUTURESTORE_ID,
34+
featurestore_id=constants.FEATURESTORE_ID,
3535
online_store_fixed_node_count=constants.ONLINE_STORE_FIXED_NODE_COUNT,
3636
sync=constants.SYNC,
3737
)

‎samples/model-builder/delete_featurestore_sample_test.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ def test_delete_featurestore_sample(
2323
delete_featurestore_sample.delete_featurestore_sample(
2424
project=constants.PROJECT,
2525
location=constants.LOCATION,
26-
featurestore_name=constants.FEAUTURESTORE_NAME,
26+
featurestore_name=constants.FEATURESTORE_NAME,
2727
sync=constants.SYNC,
2828
force=constants.FORCE,
2929
)
@@ -33,7 +33,7 @@ def test_delete_featurestore_sample(
3333
)
3434

3535
mock_get_featurestore.assert_called_once_with(
36-
featurestore_name=constants.FEAUTURESTORE_NAME
36+
featurestore_name=constants.FEATURESTORE_NAME
3737
)
3838

3939
mock_delete_featurestore.assert_called_once_with(

‎samples/model-builder/test_constants.py‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -199,8 +199,13 @@
199199
MODEL_TYPE = "CLOUD"
200200

201201
# Feature store constants
202-
FEAUTURESTORE_NAME = "projects/123/locations/us-central1/featurestores/featurestore_id"
203-
FEAUTURESTORE_ID = "featurestore_id"
202+
FEATURESTORE_ID = "featurestore_id"
203+
FEATURESTORE_NAME = "projects/123/locations/us-central1/featurestores/featurestore_id"
204+
ENTITY_TYPE_ID = "entity_type_id"
205+
ENTITY_TYPE_NAME = "projects/123/locations/us-central1/featurestores/featurestore_id/entityTypes/entity_type_id"
206+
FEATURE_ID = "feature_id"
207+
FEATURE_NAME = "projects/123/locations/us-central1/featurestores/featurestore_id/entityTypes/entity_type_id/features/feature_id"
208+
FEATURE_VALUE_TYPE = "INT64"
204209
ONLINE_STORE_FIXED_NODE_COUNT = 1
205210
SYNC = True
206211
FORCE = True

0 commit comments

Comments
 (0)