From 9a269b45d3886be77df47fec146d6cdcb9f28d1c Mon Sep 17 00:00:00 2001
From: Math Fortin <mathieu@fortin.info>
Date: Mon, 27 Nov 2023 09:18:15 -0500
Subject: [PATCH] feat: Implements placement options and license specifications
 (#607)

---------

Co-authored-by: Dan Webb <dan.webb@damacus.io>
---
 lib/kitchen/driver/aws/instance_generator.rb  | 37 ++++++++++++++++++
 .../driver/aws/instance_generator_spec.rb     | 39 ++++++++++++++++++-
 2 files changed, 75 insertions(+), 1 deletion(-)

diff --git a/lib/kitchen/driver/aws/instance_generator.rb b/lib/kitchen/driver/aws/instance_generator.rb
index 7b02b82d..c14834e8 100644
--- a/lib/kitchen/driver/aws/instance_generator.rb
+++ b/lib/kitchen/driver/aws/instance_generator.rb
@@ -189,6 +189,43 @@ def ec2_instance_data
               i[:placement] = { tenancy: }
             end
           end
+          placement = config[:placement]
+          if placement
+            unless i.key?(:placement)
+              i[:placement] = {}
+            end
+            if placement[:affinity]
+              i[:placement][:affinity] = placement[:affinity]
+            end
+            if placement[:availability_zone]
+              i[:placement][:availability_zone] = placement[:availability_zone]
+            end
+            if placement[:group_id] && !placement[:group_name]
+              i[:placement][:group_id] = placement[:group_id]
+            end
+            if placement[:group_name] && !placement[:group_id]
+              i[:placement][:group_name] = placement[:group_name]
+            end
+            if placement[:host_id]
+              i[:placement][:host_id] = placement[:host_id]
+            end
+            if placement[:host_resource_group_arn]
+              i[:placement][:host_resource_group_arn] = placement[:host_resource_group_arn]
+            end
+            if placement[:partition_number]
+              i[:placement][:partition_number] = placement[:partition_number]
+            end
+            if placement[:tenancy]
+              i[:placement][:tenancy] = placement[:tenancy]
+            end
+          end
+          license_specifications = config[:licenses]
+          if license_specifications
+            i[:licenses] = {}
+            license_specifications.each do |license_configuration_arn|
+              i[:licenses].append({ license_configuration_arn: license_configuration_arn[:license_configuration_arn] })
+            end
+          end
           unless config[:instance_initiated_shutdown_behavior].nil? ||
               config[:instance_initiated_shutdown_behavior].empty?
             i[:instance_initiated_shutdown_behavior] = config[:instance_initiated_shutdown_behavior]
diff --git a/spec/kitchen/driver/aws/instance_generator_spec.rb b/spec/kitchen/driver/aws/instance_generator_spec.rb
index 7f3b859e..f0331ee4 100644
--- a/spec/kitchen/driver/aws/instance_generator_spec.rb
+++ b/spec/kitchen/driver/aws/instance_generator_spec.rb
@@ -675,6 +675,37 @@
         }
       end
 
+      context "when placement host resource group arn and licenses are provided" do
+        let(:config) do
+          {
+            region: "eu-east-1",
+            placement: {
+              host_resource_group_arn: "arn:aws:ec2:us-east-1:123456789012:resource-group/my-group",
+            },
+            licenses: [
+              {
+                license_configuration_arn: "arn:aws:ec2:us-east-1:123456789012:license-configuration/my-license",
+              }
+            ],
+          }
+        end
+
+        it "adds the region to it in the instance data" do
+          expect(generator.ec2_instance_data).to eq(
+            instance_type: nil,
+            ebs_optimized: nil,
+            image_id: nil,
+            key_name: nil,
+            subnet_id: nil,
+            private_ip_address: nil,
+            max_count: 1,
+            min_count: 1,
+            placement: { availability_zone: "eu-east-1c",
+                         tenancy: "dedicated" }
+          )
+        end
+      end
+
       it "returns the maximum data" do
         expect(generator.ec2_instance_data).to eq(
           instance_type: "micro",
@@ -694,6 +725,9 @@
             },
           ],
           iam_instance_profile: { name: "iam-123" },
+          licenses: [{
+            license_configuration_arn: "arn:aws:ec2:us-east-1:123456789012:license-configuration/my-license",
+          }],
           network_interfaces: [{
             device_index: 0,
             associate_public_ip_address: true,
@@ -702,7 +736,10 @@
             groups: ["sg-789"],
             private_ip_address: "0.0.0.0",
           }],
-          placement: { availability_zone: "eu-west-1a" },
+          placement: {
+            availability_zone: "eu-west-1a",
+            host_resource_group_arn: "arn:aws:ec2:us-east-1:123456789012:resource-group/my-group"
+          },
           user_data: Base64.encode64("foo"),
           max_count: 1,
           min_count: 1,