|
| 1 | +"""Tests for the ``dagster-aws/<version>`` product token and the optional |
| 2 | +``user_agent_extra`` knob on ``S3Resource`` / ``S3FileManagerResource``. |
| 3 | +
|
| 4 | +Every boto3 S3 client constructed via these resources carries |
| 5 | +``dagster-aws/<version>`` in its ``user_agent_extra``, and works against AWS S3 |
| 6 | +as well as any S3-compatible backend such as Backblaze B2. Any caller-supplied |
| 7 | +``user_agent_extra`` is appended after the product token with a space |
| 8 | +separator, and never replaces it. |
| 9 | +""" |
| 10 | + |
| 11 | +from importlib.metadata import PackageNotFoundError, version |
| 12 | + |
| 13 | +from dagster_aws.s3 import S3Resource |
| 14 | +from dagster_aws.s3.resources import S3FileManagerResource |
| 15 | +from dagster_aws.s3.utils import construct_s3_client |
| 16 | + |
| 17 | +try: |
| 18 | + DAGSTER_AWS_PRODUCT = f"dagster-aws/{version('dagster-aws')}" |
| 19 | +except PackageNotFoundError: |
| 20 | + DAGSTER_AWS_PRODUCT = "dagster-aws/dev" |
| 21 | + |
| 22 | + |
| 23 | +def test_construct_s3_client_default_carries_product_token() -> None: |
| 24 | + """With no caller-supplied extra, only ``dagster-aws/<version>`` is appended.""" |
| 25 | + client = construct_s3_client(max_attempts=5) |
| 26 | + extra = client.meta.config.user_agent_extra or "" |
| 27 | + assert extra.strip() == DAGSTER_AWS_PRODUCT |
| 28 | + # The base SDK token is preserved, not replaced. |
| 29 | + assert "Boto3/" in client.meta.config.user_agent |
| 30 | + assert DAGSTER_AWS_PRODUCT in client.meta.config.user_agent |
| 31 | + |
| 32 | + |
| 33 | +def test_construct_s3_client_appends_user_agent_extra_after_product_token() -> None: |
| 34 | + """Caller-supplied tokens are appended after the product token, not in place of it.""" |
| 35 | + caller_token = "my-platform/1.2.3" |
| 36 | + client = construct_s3_client(max_attempts=5, user_agent_extra=caller_token) |
| 37 | + extra = client.meta.config.user_agent_extra or "" |
| 38 | + assert DAGSTER_AWS_PRODUCT in extra |
| 39 | + assert caller_token in extra |
| 40 | + # Order: product token first, caller token after, space-separated. |
| 41 | + assert extra == f"{DAGSTER_AWS_PRODUCT} {caller_token}" |
| 42 | + assert "Boto3/" in client.meta.config.user_agent |
| 43 | + assert caller_token in client.meta.config.user_agent |
| 44 | + |
| 45 | + |
| 46 | +def test_construct_s3_client_empty_user_agent_extra_treated_as_unset() -> None: |
| 47 | + """An empty string should not change behavior: only the product token is appended.""" |
| 48 | + client = construct_s3_client(max_attempts=5, user_agent_extra="") |
| 49 | + extra = client.meta.config.user_agent_extra or "" |
| 50 | + assert extra.strip() == DAGSTER_AWS_PRODUCT |
| 51 | + |
| 52 | + |
| 53 | +def test_construct_s3_client_preserves_retry_config() -> None: |
| 54 | + """Merging the product token must not drop the retry config already in the chain.""" |
| 55 | + client = construct_s3_client(max_attempts=7, user_agent_extra="my-platform/1.2.3") |
| 56 | + retries = client.meta.config.retries |
| 57 | + # botocore records the initial call plus retries as ``total_max_attempts``. |
| 58 | + assert retries["total_max_attempts"] == 8 |
| 59 | + assert retries["mode"] == "standard" |
| 60 | + extra = client.meta.config.user_agent_extra or "" |
| 61 | + assert DAGSTER_AWS_PRODUCT in extra |
| 62 | + assert "my-platform/1.2.3" in extra |
| 63 | + |
| 64 | + |
| 65 | +def test_s3_resource_default_carries_product_token() -> None: |
| 66 | + client = S3Resource().get_client() |
| 67 | + extra = client.meta.config.user_agent_extra or "" |
| 68 | + assert extra.strip() == DAGSTER_AWS_PRODUCT |
| 69 | + |
| 70 | + |
| 71 | +def test_s3_resource_threads_user_agent_extra_to_client() -> None: |
| 72 | + client = S3Resource(user_agent_extra="my-platform/1.2.3").get_client() |
| 73 | + extra = client.meta.config.user_agent_extra or "" |
| 74 | + assert extra == f"{DAGSTER_AWS_PRODUCT} my-platform/1.2.3" |
| 75 | + |
| 76 | + |
| 77 | +def test_s3_resource_user_agent_extra_with_b2_endpoint() -> None: |
| 78 | + """End-to-end shape a Backblaze B2 user would write.""" |
| 79 | + resource = S3Resource( |
| 80 | + endpoint_url="https://s3.us-west-004.backblazeb2.com", |
| 81 | + aws_access_key_id="00400", |
| 82 | + aws_secret_access_key="K004", |
| 83 | + user_agent_extra="my-platform/1.2.3", |
| 84 | + ) |
| 85 | + client = resource.get_client() |
| 86 | + assert client.meta.endpoint_url == "https://s3.us-west-004.backblazeb2.com" |
| 87 | + extra = client.meta.config.user_agent_extra or "" |
| 88 | + assert DAGSTER_AWS_PRODUCT in extra |
| 89 | + assert "my-platform/1.2.3" in extra |
| 90 | + |
| 91 | + |
| 92 | +def test_s3_file_manager_resource_threads_user_agent_extra() -> None: |
| 93 | + file_manager = S3FileManagerResource( |
| 94 | + s3_bucket="my-bucket", |
| 95 | + user_agent_extra="my-platform/1.2.3", |
| 96 | + ).get_client() |
| 97 | + # ``S3FileManager`` stores the wrapped boto3 client as ``_s3_session``. |
| 98 | + extra = file_manager._s3_session.meta.config.user_agent_extra or "" # noqa: SLF001 |
| 99 | + assert extra == f"{DAGSTER_AWS_PRODUCT} my-platform/1.2.3" |
| 100 | + |
| 101 | + |
| 102 | +def test_s3_file_manager_resource_default_carries_product_token() -> None: |
| 103 | + file_manager = S3FileManagerResource(s3_bucket="my-bucket").get_client() |
| 104 | + extra = file_manager._s3_session.meta.config.user_agent_extra or "" # noqa: SLF001 |
| 105 | + assert extra.strip() == DAGSTER_AWS_PRODUCT |
0 commit comments