Add EC2 runtime test for assume-role path

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()