File: test_sagemaker_feature_groups.py

package info (click to toggle)
python-moto 5.1.18-3
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 116,520 kB
  • sloc: python: 636,725; javascript: 181; makefile: 39; sh: 3
file content (118 lines) | stat: -rw-r--r-- 4,942 bytes parent folder | download | duplicates (2)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
"""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"