Skip to content

Commit 42d9c1e

Browse files
knakadajaykarpur
authored andcommitted
change: make test's custom bucket include region and account name. (#1183)
1 parent 135ffa7 commit 42d9c1e

File tree

1 file changed

+12
-5
lines changed

1 file changed

+12
-5
lines changed

tests/integ/test_processing.py

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -22,12 +22,13 @@
2222

2323
from sagemaker.processing import ProcessingInput, ProcessingOutput, ScriptProcessor, Processor
2424
from sagemaker.sklearn.processing import SKLearnProcessor
25+
from sagemaker.utils import sts_regional_endpoint
2526
from tests.integ import DATA_DIR
2627
from tests.integ.kms_utils import get_or_create_kms_key
2728

2829
ROLE = "SageMakerRole"
2930
DEFAULT_REGION = "us-west-2"
30-
CUSTOM_BUCKET_PATH = "sagemaker-custom-bucket"
31+
CUSTOM_BUCKET_PATH_PREFIX = "sagemaker-custom-bucket"
3132

3233

3334
@pytest.fixture(scope="module")
@@ -49,11 +50,17 @@ def sagemaker_session_with_custom_bucket(
4950
else None
5051
)
5152

53+
region = boto_session.region_name
54+
account = boto_session.client(
55+
"sts", region_name=region, endpoint_url=sts_regional_endpoint(region)
56+
).get_caller_identity()["Account"]
57+
custom_default_bucket = "{}-{}-{}".format(CUSTOM_BUCKET_PATH_PREFIX, region, account)
58+
5259
return Session(
5360
boto_session=boto_session,
5461
sagemaker_client=sagemaker_client,
5562
sagemaker_runtime_client=runtime_client,
56-
default_bucket=CUSTOM_BUCKET_PATH,
63+
default_bucket=custom_default_bucket,
5764
)
5865

5966

@@ -255,10 +262,10 @@ def test_sklearn_with_custom_default_bucket(
255262
job_description = sklearn_processor.latest_job.describe()
256263

257264
assert job_description["ProcessingInputs"][0]["InputName"] == "dummy_input"
258-
assert CUSTOM_BUCKET_PATH in job_description["ProcessingInputs"][0]["S3Input"]["S3Uri"]
265+
assert CUSTOM_BUCKET_PATH_PREFIX in job_description["ProcessingInputs"][0]["S3Input"]["S3Uri"]
259266

260267
assert job_description["ProcessingInputs"][1]["InputName"] == "code"
261-
assert CUSTOM_BUCKET_PATH in job_description["ProcessingInputs"][1]["S3Input"]["S3Uri"]
268+
assert CUSTOM_BUCKET_PATH_PREFIX in job_description["ProcessingInputs"][1]["S3Input"]["S3Uri"]
262269

263270
assert job_description["ProcessingJobName"].startswith("test-sklearn-with-customizations")
264271

@@ -564,7 +571,7 @@ def test_processor_with_custom_bucket(
564571
job_description = processor.latest_job.describe()
565572

566573
assert job_description["ProcessingInputs"][0]["InputName"] == "code"
567-
assert CUSTOM_BUCKET_PATH in job_description["ProcessingInputs"][0]["S3Input"]["S3Uri"]
574+
assert CUSTOM_BUCKET_PATH_PREFIX in job_description["ProcessingInputs"][0]["S3Input"]["S3Uri"]
568575

569576
assert job_description["ProcessingJobName"].startswith("test-processor")
570577

0 commit comments

Comments
 (0)