diff --git a/src/braket/aws/aws_session.py b/src/braket/aws/aws_session.py index 4abd4c15b..20fdc720a 100644 --- a/src/braket/aws/aws_session.py +++ b/src/braket/aws/aws_session.py @@ -13,6 +13,7 @@ import os import os.path +import re from typing import Any, Dict, List, NamedTuple, Optional import backoff @@ -537,8 +538,13 @@ def parse_s3_uri(s3_uri: str) -> (str, str): a valid S3 URI. """ try: - assert s3_uri.startswith("s3://") - bucket, key = s3_uri.split("/", 3)[2:] + # Object URL e.g. https://my-bucket.s3.us-west-2.amazonaws.com/my/key + # S3 URI e.g. s3://my-bucket/my/key + s3_uri_match = re.match("^https://([^./]+).[sS]3.[^/]+/(.*)$", s3_uri) or re.match( + "^[sS]3://([^./]+)/(.*)$", s3_uri + ) + assert s3_uri_match + bucket, key = s3_uri_match.groups() assert bucket and key return bucket, key except (AssertionError, ValueError): diff --git a/test/unit_tests/braket/aws/test_aws_session.py b/test/unit_tests/braket/aws/test_aws_session.py index 5927fb36b..28a98bc90 100644 --- a/test/unit_tests/braket/aws/test_aws_session.py +++ b/test/unit_tests/braket/aws/test_aws_session.py @@ -709,7 +709,7 @@ def test_is_s3_uri(string, valid): @pytest.mark.parametrize( "uri, bucket, key", - [ + ( ( "s3://bucket-name-123/key/with/multiple/dirs", "bucket-name-123", @@ -720,35 +720,35 @@ def test_is_s3_uri(string, valid): "bucket-name-123", "key-with_one.dirs", ), - pytest.param( - "http://bucket-name-123/key/with/multiple/dirs", - "bucket-name-123", - "key/with/multiple/dirs", - marks=pytest.mark.xfail(raises=ValueError, strict=True), - ), - pytest.param( - "bucket-name-123/key/with/multiple/dirs", - "bucket-name-123", - "key/with/multiple/dirs", - marks=pytest.mark.xfail(raises=ValueError, strict=True), - ), - pytest.param( - "s3://bucket-name-123/", + ( + "https://bucket-name-123.s3.us-west-2.amazonaws.com/key/with/dirs", "bucket-name-123", - "", - marks=pytest.mark.xfail(raises=ValueError, strict=True), + "key/with/dirs", ), - pytest.param( - "s3://bucket-name-123", + ( + "https://bucket-name-123.S3.us-west-2.amazonaws.com/key/with/dirs", "bucket-name-123", - "", - marks=pytest.mark.xfail(raises=ValueError, strict=True), + "key/with/dirs", ), - ], + ), ) def test_parse_s3_uri(uri, bucket, key): - parsed = AwsSession.parse_s3_uri(uri) - assert bucket, key == parsed + assert bucket, key == AwsSession.parse_s3_uri(uri) + + +@pytest.mark.parametrize( + "uri", + ( + "s3://bucket.name-123/key-with_one.dirs", + "http://bucket-name-123/key/with/multiple/dirs", + "bucket-name-123/key/with/multiple/dirs", + "s3://bucket-name-123/", + "s3://bucket-name-123", + ), +) +def test_parse_s3_uri_invalid(uri): + with pytest.raises(ValueError, match=f"Not a valid S3 uri: {uri}"): + AwsSession.parse_s3_uri(uri) @pytest.mark.parametrize(