From 91ebd96737008908e728ac12e5b50d964f5eebb7 Mon Sep 17 00:00:00 2001
From: Charles Harris <charles@mexeng.com>
Date: Mon, 23 Sep 2024 14:18:23 +0100
Subject: [PATCH 1/2] Add instance_metadata_tags option + test

---
 runner_manager/models/backend.py | 15 ++++++++++++++-
 tests/unit/backend/test_aws.py   |  5 +++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/runner_manager/models/backend.py b/runner_manager/models/backend.py
index 599d7ac2..fc284346 100644
--- a/runner_manager/models/backend.py
+++ b/runner_manager/models/backend.py
@@ -3,11 +3,16 @@
 from string import Template
 from typing import Dict, List, Literal, Optional, Sequence, TypedDict
 
-from mypy_boto3_ec2.literals import InstanceTypeType, VolumeTypeType
+from mypy_boto3_ec2.literals import (
+    InstanceMetadataTagsStateType,
+    InstanceTypeType,
+    VolumeTypeType,
+)
 from mypy_boto3_ec2.type_defs import (
     BlockDeviceMappingTypeDef,
     EbsBlockDeviceTypeDef,
     IamInstanceProfileTypeDef,
+    InstanceMetadataOptionsResponseTypeDef,
     TagSpecificationTypeDef,
     TagTypeDef,
 )
@@ -141,6 +146,7 @@ class AWSConfig(BackendConfig):
         "MaxCount": int,
         "MinCount": int,
         "IamInstanceProfile": IamInstanceProfileTypeDef,
+        "MetadataOptions": InstanceMetadataOptionsResponseTypeDef,
     },
 )
 
@@ -159,6 +165,7 @@ class AWSInstanceConfig(InstanceConfig):
     volume_type: VolumeTypeType = "gp3"
     disk_size_gb: int = 20
     iam_instance_profile_arn: str = ""
+    instance_metadata_tags: InstanceMetadataTagsStateType = "disabled"
 
     def configure_instance(self, runner: Runner) -> AwsInstance:
         """Configure instance."""
@@ -197,6 +204,11 @@ def configure_instance(self, runner: Runner) -> AwsInstance:
         iam_instance_profile: IamInstanceProfileTypeDef = IamInstanceProfileTypeDef(
             Arn=self.iam_instance_profile_arn
         )
+        instance_metadata_options: InstanceMetadataOptionsResponseTypeDef = (
+            InstanceMetadataOptionsResponseTypeDef(
+                InstanceMetadataTags=self.instance_metadata_tags,
+            )
+        )
         return AwsInstance(
             ImageId=self.image,
             InstanceType=self.instance_type,
@@ -208,6 +220,7 @@ def configure_instance(self, runner: Runner) -> AwsInstance:
             MinCount=self.min_count,
             BlockDeviceMappings=block_device_mappings,
             IamInstanceProfile=iam_instance_profile,
+            MetadataOptions=instance_metadata_options,
         )
 
 
diff --git a/tests/unit/backend/test_aws.py b/tests/unit/backend/test_aws.py
index 61310dbc..b5633411 100644
--- a/tests/unit/backend/test_aws.py
+++ b/tests/unit/backend/test_aws.py
@@ -49,6 +49,7 @@ def test_aws_instance_config(runner: Runner):
         tags={"test": "test"},
         subnet_id="i-0f9b0a3b7b3b3b3b3",
         iam_instance_profile_arn="test",
+        instance_metadata_tags="enabled",
     )
     instance: AwsInstance = instance_config.configure_instance(runner)
     assert instance["ImageId"] == instance_config.image
@@ -57,6 +58,10 @@ def test_aws_instance_config(runner: Runner):
         instance["IamInstanceProfile"]["Arn"]
         == instance_config.iam_instance_profile_arn
     )
+    assert (
+        instance["MetadataOptions"]["InstanceMetadataTags"]
+        == instance_config.instance_metadata_tags
+    )
     assert runner.name in instance["UserData"]
     tags = instance["TagSpecifications"][0]["Tags"]
     assert TagTypeDef(Key="test", Value="test") in tags

From edb5285b2e9e360f81cc389fc3532a018e6f51b1 Mon Sep 17 00:00:00 2001
From: Charles Harris <charles@mexeng.com>
Date: Mon, 23 Sep 2024 14:21:23 +0100
Subject: [PATCH 2/2] Correct TypeDef

---
 runner_manager/models/backend.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/runner_manager/models/backend.py b/runner_manager/models/backend.py
index fc284346..ce75bda7 100644
--- a/runner_manager/models/backend.py
+++ b/runner_manager/models/backend.py
@@ -12,7 +12,7 @@
     BlockDeviceMappingTypeDef,
     EbsBlockDeviceTypeDef,
     IamInstanceProfileTypeDef,
-    InstanceMetadataOptionsResponseTypeDef,
+    InstanceMetadataOptionsRequestTypeDef,
     TagSpecificationTypeDef,
     TagTypeDef,
 )
@@ -146,7 +146,7 @@ class AWSConfig(BackendConfig):
         "MaxCount": int,
         "MinCount": int,
         "IamInstanceProfile": IamInstanceProfileTypeDef,
-        "MetadataOptions": InstanceMetadataOptionsResponseTypeDef,
+        "MetadataOptions": InstanceMetadataOptionsRequestTypeDef,
     },
 )
 
@@ -204,8 +204,8 @@ def configure_instance(self, runner: Runner) -> AwsInstance:
         iam_instance_profile: IamInstanceProfileTypeDef = IamInstanceProfileTypeDef(
             Arn=self.iam_instance_profile_arn
         )
-        instance_metadata_options: InstanceMetadataOptionsResponseTypeDef = (
-            InstanceMetadataOptionsResponseTypeDef(
+        instance_metadata_options: InstanceMetadataOptionsRequestTypeDef = (
+            InstanceMetadataOptionsRequestTypeDef(
                 InstanceMetadataTags=self.instance_metadata_tags,
             )
         )