119 lines
4.8 KiB
Python
119 lines
4.8 KiB
Python
"""Unit tests for sagemaker-supported APIs."""
|
|
|
|
import re
|
|
from datetime import datetime
|
|
|
|
import boto3
|
|
import pytest
|
|
from botocore.exceptions import ClientError
|
|
|
|
from moto import mock_aws
|
|
|
|
# See our Development Tips on writing tests for hints on how to write good tests:
|
|
# http://docs.getmoto.org/en/latest/docs/contributing/development_tips/tests.html
|
|
|
|
|
|
@mock_aws
|
|
def test_create_feature_group():
|
|
client = boto3.client("sagemaker", region_name="us-east-2")
|
|
resp = client.create_feature_group(
|
|
FeatureGroupName="some-feature-group-name",
|
|
RecordIdentifierFeatureName="some_record_identifier",
|
|
EventTimeFeatureName="EventTime",
|
|
FeatureDefinitions=[
|
|
{"FeatureName": "some_feature", "FeatureType": "String"},
|
|
{"FeatureName": "EventTime", "FeatureType": "Fractional"},
|
|
{"FeatureName": "some_record_identifier", "FeatureType": "String"},
|
|
],
|
|
RoleArn="arn:aws:iam::123456789012:role/AWSFeatureStoreAccess",
|
|
OfflineStoreConfig={
|
|
"DisableGlueTableCreation": False,
|
|
"S3StorageConfig": {"S3Uri": "s3://mybucket"},
|
|
},
|
|
)
|
|
|
|
assert (
|
|
resp["FeatureGroupArn"]
|
|
== "arn:aws:sagemaker:us-east-2:123456789012:feature-group/some-feature-group-name"
|
|
)
|
|
|
|
with pytest.raises(ClientError) as raised_exception:
|
|
client.create_feature_group(
|
|
FeatureGroupName="some-feature-group-name",
|
|
RecordIdentifierFeatureName="some_record_identifier",
|
|
EventTimeFeatureName="EventTime",
|
|
FeatureDefinitions=[
|
|
{"FeatureName": "some_feature", "FeatureType": "String"},
|
|
{"FeatureName": "EventTime", "FeatureType": "Fractional"},
|
|
{"FeatureName": "some_record_identifier", "FeatureType": "String"},
|
|
],
|
|
RoleArn="arn:aws:iam::123456789012:role/AWSFeatureStoreAccess",
|
|
OfflineStoreConfig={
|
|
"DisableGlueTableCreation": False,
|
|
"S3StorageConfig": {"S3Uri": "s3://mybucket"},
|
|
},
|
|
)
|
|
|
|
assert raised_exception.value.response["Error"]["Code"] == "ResourceInUse"
|
|
assert (
|
|
raised_exception.value.response["Error"]["Message"]
|
|
== "An error occurred (ResourceInUse) when calling the CreateFeatureGroup operation: Resource Already Exists: FeatureGroup with name some-feature-group-name already exists. Choose a different name.\nInfo: Feature Group 'some-feature-group-name' already exists."
|
|
)
|
|
|
|
|
|
@mock_aws
|
|
def test_describe_feature_group():
|
|
client = boto3.client("sagemaker", region_name="us-east-2")
|
|
feature_group_name = "some-feature-group-name"
|
|
record_identifier_feature_name = "some_record_identifier"
|
|
event_time_feature_name = "EventTime"
|
|
role_arn = "arn:aws:iam::123456789012:role/AWSFeatureStoreAccess"
|
|
feature_definitions = [
|
|
{"FeatureName": "some_feature", "FeatureType": "String"},
|
|
{"FeatureName": event_time_feature_name, "FeatureType": "Fractional"},
|
|
{"FeatureName": record_identifier_feature_name, "FeatureType": "String"},
|
|
]
|
|
client.create_feature_group(
|
|
FeatureGroupName=feature_group_name,
|
|
RecordIdentifierFeatureName=record_identifier_feature_name,
|
|
EventTimeFeatureName=event_time_feature_name,
|
|
FeatureDefinitions=feature_definitions,
|
|
RoleArn=role_arn,
|
|
OfflineStoreConfig={
|
|
"DisableGlueTableCreation": False,
|
|
"S3StorageConfig": {"S3Uri": "s3://mybucket/some-folder/some-subfolder"},
|
|
},
|
|
)
|
|
resp = client.describe_feature_group(FeatureGroupName=feature_group_name)
|
|
|
|
assert resp["FeatureGroupName"] == feature_group_name
|
|
assert (
|
|
resp["FeatureGroupArn"]
|
|
== "arn:aws:sagemaker:us-east-2:123456789012:feature-group/some-feature-group-name"
|
|
)
|
|
assert resp["RecordIdentifierFeatureName"] == record_identifier_feature_name
|
|
assert resp["EventTimeFeatureName"] == event_time_feature_name
|
|
assert resp["FeatureDefinitions"] == feature_definitions
|
|
assert resp["RoleArn"] == role_arn
|
|
assert re.match(
|
|
f"^{feature_group_name.replace('-', '_')}_[0-9]+$",
|
|
resp["OfflineStoreConfig"]["DataCatalogConfig"]["TableName"],
|
|
)
|
|
assert (
|
|
resp["OfflineStoreConfig"]["DataCatalogConfig"]["Catalog"] == "AwsDataCatalog"
|
|
)
|
|
assert (
|
|
resp["OfflineStoreConfig"]["DataCatalogConfig"]["Database"]
|
|
== "sagemaker_featurestore"
|
|
)
|
|
assert (
|
|
resp["OfflineStoreConfig"]["S3StorageConfig"]["S3Uri"]
|
|
== "s3://mybucket/some-folder/some-subfolder"
|
|
)
|
|
assert re.match(
|
|
f"^s3://mybucket/some-folder/some-subfolder/123456789012/us-east-2/offline-store/{feature_group_name}-[0-9]+/data$",
|
|
resp["OfflineStoreConfig"]["S3StorageConfig"]["ResolvedOutputS3Uri"],
|
|
)
|
|
assert isinstance(resp["CreationTime"], datetime)
|
|
assert resp["FeatureGroupStatus"] == "Created"
|