def test_assume_role_china_missing_secret_key(monkeypatch, caplog):
monkeypatch.setenv("LUMIGO_LOGS_EDGE_AWS_ACCESS_KEY_ID", "111")
monkeypatch.setenv("AWS_REGION", "cn-northwest-1")
with pytest.raises(ChinaMissingEnvVar):
assume_role("111111111111", "unitteset")
assert "CRITICAL" in caplog.text
assert "LUMIGO_LOGS_EDGE_AWS_SECRET_ACCESS_KEY" in caplog.text
def test_assume_role(monkeypatch):
original_client = boto3.client
mocked_client = MagicMock(side_effect=original_client)
monkeypatch.setattr(boto3, "client", mocked_client)
assume_role("111111111111", "unitteset")
mocked_client.assert_called_with("sts")
def test_assume_role_china(monkeypatch):
monkeypatch.setenv("LUMIGO_LOGS_EDGE_AWS_ACCESS_KEY_ID", "key1")
monkeypatch.setenv("LUMIGO_LOGS_EDGE_AWS_SECRET_ACCESS_KEY", "secret1")
monkeypatch.setenv("AWS_REGION", "cn-northwest-1")
original_client = boto3.client
mocked_client = MagicMock(side_effect=original_client)
monkeypatch.setattr(boto3, "client", mocked_client)
assume_role("111111111111", "unitteset")
mocked_client.assert_called_with(
"sts",
region_name="us-west-2",
aws_access_key_id="key1",
aws_secret_access_key="secret1",
)
def get_boto_client(account_id: str, target_account_id: str):
region = get_dest_region()
if (account_id != target_account_id
and target_account_id != SELF_ACCOUNT_ID) or is_china_region():
sts_response = assume_role(target_account_id, TARGET_ENV)
return boto3.client(
"firehose",
region_name=region,
aws_access_key_id=sts_response["Credentials"]["AccessKeyId"],
aws_secret_access_key=sts_response["Credentials"]
["SecretAccessKey"],
aws_session_token=sts_response["Credentials"]["SessionToken"],
)
return boto3.client("firehose")