c6f39f51-d0b4-49d1-a213-578d9a89e6c5

Committed 20 May 2025 12:52PM UTC coverage: 86.655% (+0.02%) from 86.637%

Build # c6f39f51-d0b4-49d1-a213-578d9a89e6c5

Build Type

push

circleci

Committed by

web-flow

Commit Message

Activate new GHA pipeline on PRs (#12648)

Coverage Stats

64497 of 74430 relevant lines covered (86.65%)

0.87 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

58.47

/localstack-core/localstack/testing/aws/util.py

import functools
import os
from typing import Callable, Dict, TypeVar

import boto3
import botocore
from botocore.awsrequest import AWSPreparedRequest, AWSResponse
from botocore.client import BaseClient
from botocore.compat import HTTPHeaders
from botocore.config import Config
from botocore.exceptions import ClientError

from localstack import config
from localstack.aws.api import RequestContext
from localstack.aws.connect import (
    ClientFactory,
    ExternalAwsClientFactory,
    ExternalClientFactory,
    ServiceLevelClientFactory,
)
from localstack.aws.forwarder import create_http_request
from localstack.aws.protocol.parser import create_parser
from localstack.aws.spec import LOCALSTACK_BUILTIN_DATA_PATH, load_service
from localstack.config import is_env_true
from localstack.testing.config import (
    SECONDARY_TEST_AWS_ACCESS_KEY_ID,
    SECONDARY_TEST_AWS_PROFILE,
    SECONDARY_TEST_AWS_SECRET_ACCESS_KEY,
    SECONDARY_TEST_AWS_SESSION_TOKEN,
    TEST_AWS_ACCESS_KEY_ID,
    TEST_AWS_REGION_NAME,
    TEST_AWS_SECRET_ACCESS_KEY,
)
from localstack.utils.aws.arns import get_partition
from localstack.utils.aws.request_context import get_account_id_from_request
from localstack.utils.sync import poll_condition


def is_aws_cloud() -> bool:
    return os.environ.get("TEST_TARGET", "") == "AWS_CLOUD"


def in_default_partition() -> bool:
    return is_aws_cloud() or get_partition(TEST_AWS_REGION_NAME) == "aws"


def get_lambda_logs(func_name, logs_client):
    log_group_name = f"/aws/lambda/{func_name}"
    streams = logs_client.describe_log_streams(logGroupName=log_group_name)["logStreams"]
    streams = sorted(streams, key=lambda x: x["creationTime"], reverse=True)
    log_events = logs_client.get_log_events(
        logGroupName=log_group_name, logStreamName=streams[0]["logStreamName"]
    )["events"]
    return log_events


def bucket_exists(client, bucket_name: str) -> bool:
    buckets = client.list_buckets()
    for bucket in buckets["Buckets"]:
        if bucket["Name"] == bucket_name:
            return True
    return False


def wait_for_user(keys, region_name: str):
    sts_client = create_client_with_keys(service="sts", keys=keys, region_name=region_name)

    def is_user_ready():
        try:
            sts_client.get_caller_identity()
            return True
        except ClientError as e:
            if e.response["Error"]["Code"] == "InvalidClientTokenId":
                return False
            return True

    # wait until the given user is ready, takes AWS IAM a while...
    poll_condition(is_user_ready, interval=5, timeout=20)


def create_client_with_keys(
    service: str,
    keys: Dict[str, str],
    region_name: str,
    client_config: Config = None,
):
    """
    Create a boto client with the given access key, targeted against LS per default, but to AWS if TEST_TARGET is set
    accordingly.

    :param service: Service to create the Client for
    :param keys: Access Keys
    :param region_name: Region for the client
    :param client_config:
    :return:
    """
    return boto3.client(
        service,
        region_name=region_name,
        aws_access_key_id=keys["AccessKeyId"],
        aws_secret_access_key=keys["SecretAccessKey"],
        aws_session_token=keys.get("SessionToken"),
        config=client_config,
        endpoint_url=config.internal_service_url() if not is_aws_cloud() else None,
    )


def create_request_context(
    service_name: str, operation_name: str, region: str, aws_request: AWSPreparedRequest
) -> RequestContext:
    if hasattr(aws_request.body, "read"):
        aws_request.body = aws_request.body.read()
    request = create_http_request(aws_request)

    context = RequestContext(request=request)
    context.service = load_service(service_name)
    context.operation = context.service.operation_model(operation_name=operation_name)
    context.region = region
    parser = create_parser(context.service)
    _, instance = parser.parse(context.request)
    context.service_request = instance
    context.account_id = get_account_id_from_request(context.request)
    return context


class _RequestContextClient:
    _client: BaseClient

    def __init__(self, client: BaseClient):
        self._client = client

    def __getattr__(self, item):
        target = getattr(self._client, item)
        if not isinstance(target, Callable):
            return target

        @functools.wraps(target)
        def wrapper_method(*args, **kwargs):
            service_name = self._client.meta.service_model.service_name
            operation_name = self._client.meta.method_to_api_mapping[item]
            region = self._client.meta.region_name
            prepared_request = None

            def event_handler(request: AWSPreparedRequest, **_):
                nonlocal prepared_request
                prepared_request = request
                # we need to return an AWS Response here
                aws_response = AWSResponse(
                    url=request.url, status_code=200, headers=HTTPHeaders(), raw=None
                )
                aws_response._content = b""
                return aws_response

            self._client.meta.events.register(
                f"before-send.{service_name}.{operation_name}", handler=event_handler
            )
            try:
                target(*args, **kwargs)
            except Exception:
                pass
            self._client.meta.events.unregister(
                f"before-send.{service_name}.{operation_name}", handler=event_handler
            )

            return create_request_context(
                service_name=service_name,
                operation_name=operation_name,
                region=region,
                aws_request=prepared_request,
            )

        return wrapper_method


T = TypeVar("T", bound=BaseClient)


def RequestContextClient(client: T) -> T:
    return _RequestContextClient(client)  # noqa


# Used for the aws_session, aws_client_factory and aws_client pytest fixtures
# Supports test executions against both LocalStack and production AWS

# TODO: Add the ability to use config profiles for primary and secondary clients
# See https://boto3.amazonaws.com/v1/documentation/api/latest/guide/configuration.html#using-a-configuration-file


def base_aws_session() -> boto3.Session:
    # When running against AWS, initial credentials must be read from environment or config file
    if is_aws_cloud():
        return boto3.Session()

    # Otherwise, when running against LS, use primary test credentials to start with
    # This set here in the session so that both `aws_client` and `aws_client_factory` can work without explicit creds.
    session = boto3.Session(
        aws_access_key_id=TEST_AWS_ACCESS_KEY_ID,
        aws_secret_access_key=TEST_AWS_SECRET_ACCESS_KEY,
    )
    # make sure we consider our custom data paths for legacy specs (like SQS query protocol)
    session._loader.search_paths.insert(0, LOCALSTACK_BUILTIN_DATA_PATH)
    return session


def secondary_aws_session() -> boto3.Session:
    if is_aws_cloud() and SECONDARY_TEST_AWS_PROFILE:
        return boto3.Session(profile_name=SECONDARY_TEST_AWS_PROFILE)

    # Otherwise, when running against LS or AWS, but have no profile set for the secondary account,
    # we use secondary test credentials to initialize the session.
    # This set here in the session so that both `secondary_aws_client` and `secondary_aws_client_factory` can work
    # without explicit creds.
    session = boto3.Session(
        aws_access_key_id=SECONDARY_TEST_AWS_ACCESS_KEY_ID,
        aws_secret_access_key=SECONDARY_TEST_AWS_SECRET_ACCESS_KEY,
        aws_session_token=SECONDARY_TEST_AWS_SESSION_TOKEN,
    )
    if not is_aws_cloud():
        # make sure we consider our custom data paths for legacy specs (like SQS query protocol), only if we run against
        # LocalStack
        session._loader.search_paths.append(LOCALSTACK_BUILTIN_DATA_PATH)
    return session


def base_aws_client_factory(session: boto3.Session) -> ClientFactory:
    config = None
    if is_env_true("TEST_DISABLE_RETRIES_AND_TIMEOUTS"):
        config = botocore.config.Config(
            connect_timeout=1_000,
            read_timeout=1_000,
            retries={"total_max_attempts": 1},
        )

    if is_aws_cloud():
        return ExternalAwsClientFactory(session=session, config=config)
    else:
        if not config:
            config = botocore.config.Config()

        # Prevent this fixture from using the region configured in system config
        config = config.merge(botocore.config.Config(region_name=TEST_AWS_REGION_NAME))
        return ExternalClientFactory(session=session, config=config)


def base_testing_aws_client(client_factory: ClientFactory) -> ServiceLevelClientFactory:
    # Primary test credentials are already set in the boto3 session, so they're not set here again
    return client_factory()

1	import functools	1✔
2	import os	1✔
3	from typing import Callable, Dict, TypeVar	1✔
4
5	import boto3	1✔
6	import botocore	1✔
7	from botocore.awsrequest import AWSPreparedRequest, AWSResponse	1✔
8	from botocore.client import BaseClient	1✔
9	from botocore.compat import HTTPHeaders	1✔
10	from botocore.config import Config	1✔
11	from botocore.exceptions import ClientError	1✔
12
13	from localstack import config	1✔
14	from localstack.aws.api import RequestContext	1✔
15	from localstack.aws.connect import (	1✔
16	ClientFactory,
17	ExternalAwsClientFactory,
18	ExternalClientFactory,
19	ServiceLevelClientFactory,
20	)
21	from localstack.aws.forwarder import create_http_request	1✔
22	from localstack.aws.protocol.parser import create_parser	1✔
23	from localstack.aws.spec import LOCALSTACK_BUILTIN_DATA_PATH, load_service	1✔
24	from localstack.config import is_env_true	1✔
25	from localstack.testing.config import (	1✔
26	SECONDARY_TEST_AWS_ACCESS_KEY_ID,
27	SECONDARY_TEST_AWS_PROFILE,
28	SECONDARY_TEST_AWS_SECRET_ACCESS_KEY,
29	SECONDARY_TEST_AWS_SESSION_TOKEN,
30	TEST_AWS_ACCESS_KEY_ID,
31	TEST_AWS_REGION_NAME,
32	TEST_AWS_SECRET_ACCESS_KEY,
33	)
34	from localstack.utils.aws.arns import get_partition	1✔
35	from localstack.utils.aws.request_context import get_account_id_from_request	1✔
36	from localstack.utils.sync import poll_condition	1✔
37
38
39	def is_aws_cloud() -> bool:	1✔
40	return os.environ.get("TEST_TARGET", "") == "AWS_CLOUD"	1✔
41
42
43	def in_default_partition() -> bool:	1✔
44	return is_aws_cloud() or get_partition(TEST_AWS_REGION_NAME) == "aws"	1✔
45
46
47	def get_lambda_logs(func_name, logs_client):	1✔
48	log_group_name = f"/aws/lambda/{func_name}"	1✔
49	streams = logs_client.describe_log_streams(logGroupName=log_group_name)["logStreams"]	1✔
50	streams = sorted(streams, key=lambda x: x["creationTime"], reverse=True)	1✔
51	log_events = logs_client.get_log_events(	1✔
52	logGroupName=log_group_name, logStreamName=streams[0]["logStreamName"]
53	)["events"]
54	return log_events	1✔
55
56
57	def bucket_exists(client, bucket_name: str) -> bool:	1✔
58	buckets = client.list_buckets()	×
59	for bucket in buckets["Buckets"]:	×
60	if bucket["Name"] == bucket_name:	×
61	return True	×
62	return False	×
63
64
65	def wait_for_user(keys, region_name: str):	1✔
66	sts_client = create_client_with_keys(service="sts", keys=keys, region_name=region_name)	1✔
67
68	def is_user_ready():	1✔
69	try:	1✔
70	sts_client.get_caller_identity()	1✔
71	return True	1✔
72	except ClientError as e:	×
73	if e.response["Error"]["Code"] == "InvalidClientTokenId":	×
74	return False	×
75	return True	×
76
77	# wait until the given user is ready, takes AWS IAM a while...
78	poll_condition(is_user_ready, interval=5, timeout=20)	1✔
79
80
81	def create_client_with_keys(	1✔
82	service: str,
83	keys: Dict[str, str],
84	region_name: str,
85	client_config: Config = None,
86	):
87	"""
88	Create a boto client with the given access key, targeted against LS per default, but to AWS if TEST_TARGET is set
89	accordingly.
90
91	:param service: Service to create the Client for
92	:param keys: Access Keys
93	:param region_name: Region for the client
94	:param client_config:
95	:return:
96	"""
97	return boto3.client(	1✔
98	service,
99	region_name=region_name,
100	aws_access_key_id=keys["AccessKeyId"],
101	aws_secret_access_key=keys["SecretAccessKey"],
102	aws_session_token=keys.get("SessionToken"),
103	config=client_config,
104	endpoint_url=config.internal_service_url() if not is_aws_cloud() else None,
105	)
106
107
108	def create_request_context(	1✔
109	service_name: str, operation_name: str, region: str, aws_request: AWSPreparedRequest
110	) -> RequestContext:
111	if hasattr(aws_request.body, "read"):	×
112	aws_request.body = aws_request.body.read()	×
113	request = create_http_request(aws_request)	×
114
115	context = RequestContext(request=request)	×
116	context.service = load_service(service_name)	×
117	context.operation = context.service.operation_model(operation_name=operation_name)	×
118	context.region = region	×
119	parser = create_parser(context.service)	×
120	_, instance = parser.parse(context.request)	×
121	context.service_request = instance	×
122	context.account_id = get_account_id_from_request(context.request)	×
123	return context	×
124
125
126	class _RequestContextClient:	1✔
127	_client: BaseClient	1✔
128
129	def __init__(self, client: BaseClient):	1✔
130	self._client = client	×
131
132	def __getattr__(self, item):	1✔
133	target = getattr(self._client, item)	×
134	if not isinstance(target, Callable):	×
135	return target	×
136
137	@functools.wraps(target)	×
138	def wrapper_method(args, *kwargs):	×
139	service_name = self._client.meta.service_model.service_name	×
140	operation_name = self._client.meta.method_to_api_mapping[item]	×
141	region = self._client.meta.region_name	×
142	prepared_request = None	×
143
144	def event_handler(request: AWSPreparedRequest, **_):	×
145	nonlocal prepared_request
146	prepared_request = request	×
147	# we need to return an AWS Response here
148	aws_response = AWSResponse(	×
149	url=request.url, status_code=200, headers=HTTPHeaders(), raw=None
150	)
151	aws_response._content = b""	×
152	return aws_response	×
153
154	self._client.meta.events.register(	×
155	f"before-send.{service_name}.{operation_name}", handler=event_handler
156	)
157	try:	×
158	target(args, *kwargs)	×
159	except Exception:	×
160	pass	×
161	self._client.meta.events.unregister(	×
162	f"before-send.{service_name}.{operation_name}", handler=event_handler
163	)
164
165	return create_request_context(	×
166	service_name=service_name,
167	operation_name=operation_name,
168	region=region,
169	aws_request=prepared_request,
170	)
171
172	return wrapper_method	×
173
174
175	T = TypeVar("T", bound=BaseClient)	1✔
176
177
178	def RequestContextClient(client: T) -> T:	1✔
179	return _RequestContextClient(client) # noqa	×
180
181
182	# Used for the aws_session, aws_client_factory and aws_client pytest fixtures
183	# Supports test executions against both LocalStack and production AWS
184
185	# TODO: Add the ability to use config profiles for primary and secondary clients
186	# See https://boto3.amazonaws.com/v1/documentation/api/latest/guide/configuration.html#using-a-configuration-file
187
188
189	def base_aws_session() -> boto3.Session:	1✔
190	# When running against AWS, initial credentials must be read from environment or config file
191	if is_aws_cloud():	1✔
192	return boto3.Session()	×
193
194	# Otherwise, when running against LS, use primary test credentials to start with
195	# This set here in the session so that both `aws_client` and `aws_client_factory` can work without explicit creds.
196	session = boto3.Session(	1✔
197	aws_access_key_id=TEST_AWS_ACCESS_KEY_ID,
198	aws_secret_access_key=TEST_AWS_SECRET_ACCESS_KEY,
199	)
200	# make sure we consider our custom data paths for legacy specs (like SQS query protocol)
201	session._loader.search_paths.insert(0, LOCALSTACK_BUILTIN_DATA_PATH)	1✔
202	return session	1✔
203
204
205	def secondary_aws_session() -> boto3.Session:	1✔
206	if is_aws_cloud() and SECONDARY_TEST_AWS_PROFILE:	1✔
207	return boto3.Session(profile_name=SECONDARY_TEST_AWS_PROFILE)	×
208
209	# Otherwise, when running against LS or AWS, but have no profile set for the secondary account,
210	# we use secondary test credentials to initialize the session.
211	# This set here in the session so that both `secondary_aws_client` and `secondary_aws_client_factory` can work
212	# without explicit creds.
213	session = boto3.Session(	1✔
214	aws_access_key_id=SECONDARY_TEST_AWS_ACCESS_KEY_ID,
215	aws_secret_access_key=SECONDARY_TEST_AWS_SECRET_ACCESS_KEY,
216	aws_session_token=SECONDARY_TEST_AWS_SESSION_TOKEN,
217	)
218	if not is_aws_cloud():	1✔
219	# make sure we consider our custom data paths for legacy specs (like SQS query protocol), only if we run against
220	# LocalStack
221	session._loader.search_paths.append(LOCALSTACK_BUILTIN_DATA_PATH)	1✔
222	return session	1✔
223
224
225	def base_aws_client_factory(session: boto3.Session) -> ClientFactory:	1✔
226	config = None	1✔
227	if is_env_true("TEST_DISABLE_RETRIES_AND_TIMEOUTS"):	1✔
228	config = botocore.config.Config(	×
229	connect_timeout=1_000,
230	read_timeout=1_000,
231	retries={"total_max_attempts": 1},
232	)
233
234	if is_aws_cloud():	1✔
235	return ExternalAwsClientFactory(session=session, config=config)	×
236	else:
237	if not config:	1✔
238	config = botocore.config.Config()	1✔
239
240	# Prevent this fixture from using the region configured in system config
241	config = config.merge(botocore.config.Config(region_name=TEST_AWS_REGION_NAME))	1✔
242	return ExternalClientFactory(session=session, config=config)	1✔
243
244
245	def base_testing_aws_client(client_factory: ClientFactory) -> ServiceLevelClientFactory:	1✔
246	# Primary test credentials are already set in the boto3 session, so they're not set here again
247	return client_factory()	1✔

localstack / localstack / c6f39f51-d0b4-49d1-a213-578d9a89e6c5

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous