How to use get_aws_account_id method in localstack

Best Python code snippet using localstack_python

test_client.py

Source:test_client.py

...124async def test_async_sns_client(async_sns_client):125    assert async_sns_client._session126    assert async_sns_client.aws_account_id is None127@pytest.mark.asyncio128async def test_async_sns_client_get_aws_account_id():129    async_sns_client = AsyncSNSClient()130    async_sns_client._session.create_client = mock.Mock()131    mock_get_caller_identity = asynctest.CoroutineMock()132    async_sns_client._session.create_client.return_value.get_caller_identity = mock_get_caller_identity133    assert await async_sns_client.get_aws_account_id()134    async_sns_client._session.create_client.assert_called_once_with(135        "sts", endpoint_url=async_sns_client.endpoint_url, use_ssl=async_sns_client.use_ssl136    )137    mock_get_caller_identity.assert_called_once_with()138@pytest.mark.asyncio139async def test_async_sns_client_get_aws_account_id_custom_endpoint(monkeypatch):140    monkeypatch.delenv("AWS_ACCOUNT_ID", raising=False)141    async_sns_client = AsyncSNSClient(endpoint_url="http://hue.com/")142    result = await async_sns_client.get_aws_account_id()143    assert result == "000000000000"144@pytest.mark.asyncio145async def test_async_sns_client_get_aws_account_id_custom_endpoint_custom_account_id(146    async_sns_client, account_id, monkeypatch147):148    monkeypatch.setenv("AWS_ACCOUNT_ID", account_id)149    async_sns_client = AsyncSNSClient(endpoint_url="http://hue.com/")150    result = await async_sns_client.get_aws_account_id()151    assert result == account_id152@pytest.mark.asyncio153async def test_async_sns_client_get_topic_arn(async_sns_client, account_id):154    topic_name = "top__ico"155    async_sns_client.get_aws_account_id = asynctest.CoroutineMock(return_value=account_id)156    async_sns_client.region_name = "us-east-1"157    result = await async_sns_client.get_topic_arn(topic_name)158    assert result == "arn:aws:sns:us-east-1:{}:{}".format(account_id, topic_name)159    async_sns_client.get_aws_account_id.assert_called_once()160@pytest.mark.asyncio161async def test_async_sns_client_get_topic_arn_cached_account_id(async_sns_client):162    topic_name = "top__ico"163    async_sns_client.aws_account_id = "000000000000"164    async_sns_client.get_aws_account_id = asynctest.CoroutineMock()...

createrole.py

Source:createrole.py

...33def run_command(cmd, **kwargs):34    "Passes args to subprocess.run, printing out the command before running it"35    print("Running:", " ".join(cmd))36    return subprocess.run(cmd, **kwargs)37def get_aws_account_id():38    """Account ID for the current aws cli configuration, caching where possible"""39    global AWS_ACCOUNT_ID40    if not AWS_ACCOUNT_ID:41        run = run_command(42            ["aws", "sts", "get-caller-identity"],43            check=True,44            capture_output=True,45            text=True,46        )47        AWS_ACCOUNT_ID = json.loads(run.stdout)["Account"]48    return AWS_ACCOUNT_ID49def create_aws_role(rolename):50    """Create an AWS role of the given rolename, with the current account as the principal"""51    userarns = [52        f"arn:aws:iam::{get_aws_account_id()}:user/{username}"53        for username in ARGS.usernames54    ]55    ROLE_TEMPLATE["Statement"][0]["Principal"]["AWS"] = userarns56    rolef = ARGS.workingdir / "role_policy.json"57    json.dump(ROLE_TEMPLATE, rolef.open("w"))58    run_command(59        [60            "aws",61            "iam",62            "create-role",63            "--role-name",64            rolename,65            "--assume-role-policy-document",66            f"file://{str(rolef)}",67        ],68    )69def XXXcreate_aws_policy(policyname, rolename):70    "Create <policyname> to allow <rolename> to be assumed"71    POLICY_TEMPLATE["Statement"][0][72        "Resource"73    ] = f"arn:aws:iam::{get_aws_account_id()}:role/{rolename}"74    policyf = ARGS.workingdir / "assume_role_policy.json"75    json.dump(POLICY_TEMPLATE, policyf.open("w"))76    print(policyf)77    # run_command(78    #     [79    #         "aws",80    #         "iam",81    #         "create-policy",82    #         "--policy-name",83    #         policyname,84    #         "--policy-document",85    #         f"file://{str(policyf)}",86    #     ]87    # )88def create_aws_group(groupname, rolename, policyname):89    """Create <groupname>, and give it the policy <policyname>, allowing assume role.90    Users are then assigned to the group.91    """92    run_command(["aws", "iam", "create-group", "--group-name", groupname])93    POLICY_TEMPLATE["Statement"][0][94        "Resource"95    ] = f"arn:aws:iam::{get_aws_account_id()}:role/{rolename}"96    policyf = ARGS.workingdir / "assume_role_policy.json"97    json.dump(POLICY_TEMPLATE, policyf.open("w"))98    run_command(99        [100            "aws",101            "iam",102            "put-group-policy",103            "--group-name",104            groupname,105            "--policy-name",106            policyname,107            "--policy-document",108            f"file://{str(policyf)}",109        ]110    )111    # Assign users112    for username in ARGS.usernames:113        run_command(114            [115                "aws",116                "iam",117                "add-user-to-group",118                "--group-name",119                groupname,120                "--user-name",121                username,122            ]123        )124def create_identity_mapping(rolename, eksusername):125    """Creates an identity mapping from k8s to the AWS IAM"""126    run_command(127        [128            "eksctl",129            "create",130            "iamidentitymapping",131            "--cluster",132            ARGS.clustername,133            "--arn",134            f"arn:aws:iam::{get_aws_account_id()}:role/{rolename}",135            "--username",136            eksusername,137        ]138    )139def getargs():140    parser = argparse.ArgumentParser(141        description="Create a AWS role and group, and link to k8s. The currently "142        "active aws cli and kubectl context is used. Names and roles are prefixes "143        'with "eks-" and finish with their purpose'144    )145    parser.add_argument(146        "accountname",147        help="The base name of the account/role/group to create",148    )...

aws_athena.py

Source:aws_athena.py

2import time3from string import Template4# Global holding the AWS account ID this is executing in5account_id = 06def get_aws_account_id(session):7    global account_id8    if account_id == 0:9        account_id = session.client('sts').get_caller_identity()['Account']10    return account_id11# Creates a table in Athena for a load balancer12# Uses template files as the fields are different for ALB/ELB13def create_athena_elb_table(force, database, elb_type, bucket, service_name, session):14    status = True15    # s3://BUCKET/ALB PREFIX/AWSLogs/ACCOUNT NUMBER/elasticloadbalancing/us-east-116    s3_location = 's3://' + bucket + '/' + service_name + \17        '/AWSLogs/' + \18        str(get_aws_account_id(session)) + \19        '/elasticloadbalancing/' + \20        session.region_name21    # Tables cannot have dashes22    table_name = service_name.lower().replace('-', '_')23    # Templates are slightly different for the 2 types24    if elb_type == "ALB":25        filein = open('sql/alb.sql')26    elif elb_type == "ELB":27        filein = open('sql/elb.sql')28    if force:29        # We need to remove the table first30        print("Force option specified - recreating Athena table")31        remove_table_sql = "DROP TABLE " + table_name32        query_id = submit_query(remove_table_sql, database, session)33        if query_id:34            if wait_for_query_to_complete(query_id, session):35                status = True36            else:37                status = False38    if filein and status:39        src = Template(filein.read())40        vals = {'table_name': table_name, 's3_location': s3_location}41        create_table_query = src.safe_substitute(vals)42        query_id = submit_query(create_table_query, database, session)43        if query_id:44            if wait_for_query_to_complete(query_id, session):45                status = True46            else:47                status = False48        else:49            status = False50    return status51def create_athena_database(db_name, session):52    status = False53    query_string = "CREATE DATABASE IF NOT EXISTS " + db_name54    query_id = submit_query(query_string, 'default', session)55    if query_id:56        if wait_for_query_to_complete(query_id, session):57            status = True58        else:59            status = False60    else:61        status = False62    return status63def submit_query(query, database, session):64    output_location = 's3://aws-athena-query-results-' + str(get_aws_account_id(session)) + "-" + session.region_name65    query_id = None66    response = None67    client = session.client('athena')68    try:69        response = client.start_query_execution(70            QueryString=query,71            QueryExecutionContext={72                'Database': database73            },74            ResultConfiguration={75                'OutputLocation': output_location,76                'EncryptionConfiguration': {77                    'EncryptionOption': 'SSE_S3'78                }...

Automation Testing Tutorials

Learn to execute automation testing from scratch with LambdaTest Learning Hub. Right from setting up the prerequisites to run your first automation test, to following best practices and diving deeper into advanced test scenarios. LambdaTest Learning Hubs compile a list of step-by-step guides to help you be proficient with different test automation frameworks i.e. Selenium, Cypress, TestNG etc.