diff --git a/agent/nix_builder_autoscaler/tests/test_runtime_ec2.py b/agent/nix_builder_autoscaler/tests/test_runtime_ec2.py index de8b26a..765288d 100644 --- a/agent/nix_builder_autoscaler/tests/test_runtime_ec2.py +++ b/agent/nix_builder_autoscaler/tests/test_runtime_ec2.py @@ -1,6 +1,6 @@ """Unit tests for the EC2 runtime adapter using botocore Stubber.""" -from datetime import UTC, datetime +from datetime import UTC, datetime, timedelta from unittest.mock import patch import boto3 @@ -462,3 +462,62 @@ class TestErrorClassification: with pytest.raises(RuntimeAdapterError) as exc_info: runtime.launch_instance("slot001", "#!/bin/bash") assert exc_info.value.category == "throttled" + + +class TestAssumeRole: + def test_uses_assumed_role_credentials_for_ec2_calls(self): + config = _make_config() + config.assume_role_arn = "arn:aws:iam::210987654321:role/buildbot-autoscaler-controller" + + base_ec2 = boto3.client("ec2", region_name="us-east-1") + assumed_ec2 = boto3.client("ec2", region_name="us-east-1") + sts_client = boto3.client("sts", region_name="us-east-1") + + sts_stubber = Stubber(sts_client) + sts_stubber.add_response( + "assume_role", + { + "Credentials": { + "AccessKeyId": "ASIAAAAAAAAAAAAAAAAA", + "SecretAccessKey": "s" * 40, + "SessionToken": "t" * 256, + "Expiration": datetime.now(UTC) + timedelta(hours=1), + }, + "AssumedRoleUser": { + "AssumedRoleId": "AROA1234567890EXAMPLE:nix-builder-autoscaler", + "Arn": ( + "arn:aws:sts::210987654321:assumed-role/" + "buildbot-autoscaler-controller/nix-builder-autoscaler" + ), + }, + }, + { + "RoleArn": config.assume_role_arn, + "RoleSessionName": "nix-builder-autoscaler", + }, + ) + sts_stubber.activate() + + assumed_stubber = Stubber(assumed_ec2) + assumed_stubber.add_response( + "run_instances", + {"Instances": [{"InstanceId": "i-assumed"}], "OwnerId": "210987654321"}, + ) + assumed_stubber.activate() + + real_boto3_client = boto3.client + + def _patched_client(service_name, **kwargs): + if service_name == "sts": + return sts_client + if service_name == "ec2" and kwargs.get("aws_access_key_id") == "ASIAAAAAAAAAAAAAAAAA": + return assumed_ec2 + return real_boto3_client(service_name, **kwargs) + + with patch("nix_builder_autoscaler.runtime.ec2.boto3.client", side_effect=_patched_client): + runtime = EC2Runtime(config, _client=base_ec2) + instance_id = runtime.launch_instance("slot001", "#!/bin/bash") + + assert instance_id == "i-assumed" + sts_stubber.assert_no_pending_responses() + assumed_stubber.assert_no_pending_responses()