Add find_ami_id logic instead of pass argument --instance-id

integ/src/ec2_provider.rs

@@ -40,18 +40,25 @@ impl Ec2Creator {}
 
 pub async fn create_ec2_instance(
     cluster: ClusterInfo,
-    node_ami: &str,
+    ami_arch: &str,
+    bottlerocket_version: &str,
 ) -> ProviderResult<CreatedEc2Instances> {
     // Setup aws_sdk_config and clients.
     let region_provider = RegionProviderChain::first_try(Some(Region::new(cluster.region.clone())));
     let shared_config = aws_config::from_env().region(region_provider).load().await;
     let ec2_client = aws_sdk_ec2::Client::new(&shared_config);
+    let ssm_client = aws_sdk_ssm::Client::new(&shared_config);
 
     // Prepare security groups
     let mut security_groups = vec![];
     security_groups.append(&mut cluster.nodegroup_sg.clone());
     security_groups.append(&mut cluster.clustershared_sg.clone());
 
+    // Prepare ami id
+    //default eks_version to the version that matches cluster
+    let eks_version = cluster.version;
+    let node_ami = find_ami_id(&ssm_client, ami_arch, bottlerocket_version, &eks_version).await?;
+
     // Prepare instance type
     let instance_type = instance_type(&ec2_client, &node_ami).await?;
 
@@ -138,6 +145,30 @@ pub async fn terminate_ec2_instance(cluster: ClusterInfo) -> ProviderResult<()>
     Ok(())
 }
 
+// Find the node ami id to use.
+async fn find_ami_id(
+    ssm_client: &aws_sdk_ssm::Client,
+    arch: &str,
+    br_version: &str,
+    eks_version: &str,
+) -> ProviderResult<String> {
+    let parameter_name = format!(
+        "/aws/service/bottlerocket/aws-k8s-{}/{}/{}/image_id",
+        eks_version, arch, br_version
+    );
+    let ami_id = ssm_client
+        .get_parameter()
+        .name(parameter_name)
+        .send()
+        .await
+        .context("Unable to get ami id")?
+        .parameter
+        .context("Unable to get ami id")?
+        .value
+        .context("ami id is missing")?;
+    Ok(ami_id)
+}
+
 /// Determine the instance type to use. If provided use that one. Otherwise, for `x86_64` use `m5.large`
 /// and for `aarch64` use `m6g.large`
 async fn instance_type(ec2_client: &aws_sdk_ec2::Client, node_ami: &str) -> ProviderResult<String> {

integ/src/eks_provider.rs

@@ -16,6 +16,7 @@ use std::process::Command;
 pub struct ClusterInfo {
     pub name: String,
     pub region: String,
+    pub version: String,
     pub endpoint: String,
     pub certificate: String,
     pub public_subnet_ids: Vec<String>,
@@ -60,6 +61,7 @@ pub async fn get_cluster_info(cluster_name: &str, region: &str) -> ProviderResul
     let ec2_client = aws_sdk_ec2::Client::new(&shared_config);
     let iam_client = aws_sdk_iam::Client::new(&shared_config);
 
+    let eks_version = eks_version(&eks_client, cluster_name).await?;
     let eks_subnet_ids = eks_subnet_ids(&eks_client, cluster_name).await?;
     let endpoint = endpoint(&eks_client, cluster_name).await?;
     let certificate = certificate(&eks_client, cluster_name).await?;
@@ -113,6 +115,7 @@ pub async fn get_cluster_info(cluster_name: &str, region: &str) -> ProviderResul
     Ok(ClusterInfo {
         name: cluster_name.to_string(),
         region: region.to_string(),
+        version: eks_version,
         endpoint,
         certificate,
         public_subnet_ids,
@@ -124,6 +127,28 @@ pub async fn get_cluster_info(cluster_name: &str, region: &str) -> ProviderResul
     })
 }
 
+async fn eks_version(
+    eks_client: &aws_sdk_eks::Client,
+    cluster_name: &str,
+) -> ProviderResult<String> {
+    let describe_results = eks_client
+        .describe_cluster()
+        .name(cluster_name)
+        .send()
+        .await
+        .context("Unable to get eks describe cluster")?;
+
+    // Extract the eks version from the cluster.
+    describe_results
+        .cluster
+        .as_ref()
+        .context("Response missing cluster field")?
+        .version
+        .as_ref()
+        .context("Cluster missing version field")
+        .map(|ids| ids.clone())
+}
+
 async fn eks_subnet_ids(
     eks_client: &aws_sdk_eks::Client,
     cluster_name: &str,

integ/src/main.rs

@@ -1,11 +1,14 @@
+use lazy_static::lazy_static;
 use log::info;
-use snafu::{OptionExt, ResultExt};
+use snafu::{ensure, OptionExt, ResultExt};
 use std::env::temp_dir;
 use std::fs;
 use std::process;
 use structopt::StructOpt;
 use tokio::time::{sleep, Duration};
 
+use aws_sdk_ec2::model::ArchitectureValues;
+
 use integ::ec2_provider::{create_ec2_instance, terminate_ec2_instance};
 use integ::eks_provider::{get_cluster_info, write_kubeconfig};
 use integ::error::ProviderError;
@@ -20,9 +23,17 @@ const DEFAULT_KUBECONFIG_FILE_NAME: &str = "kubeconfig.yaml";
 const DEFAULT_REGION: &str = "us-west-2";
 const CLUSTER_NAME: &str = "brupop-integration-test";
 
+//The default values for AMI ID
+const AMI_ARCH: &str = "x86_64";
+
 /// This value configure how long it sleeps between create instance and label instance.
 const INTEGRATION_TEST_DELAY: Duration = Duration::from_secs(60);
 
+lazy_static! {
+    static ref ARCHES: Vec<ArchitectureValues> =
+        vec![ArchitectureValues::Arm64, ArchitectureValues::X8664];
+}
+
 #[tokio::main]
 async fn main() {
     env_logger::init();
@@ -60,9 +71,12 @@ enum SubCommand {
 
 // Stores user-supplied arguments for the 'integration-test' subcommand.
 #[derive(StructOpt, Debug)]
-struct IntegrationTestArgs {
-    #[structopt(long = "--instance-ami-id")]
-    instance_ami_id: String,
+pub struct IntegrationTestArgs {
+    #[structopt(long = "--bottlerocket-version")]
+    bottlerocket_version: String,
+
+    #[structopt(long = "--arch", default_value = AMI_ARCH)]
+    ami_arch: String,
 }
 
 async fn generate_kubeconfig(arguments: &Arguments) -> Result<String> {
@@ -110,10 +124,30 @@ fn get_kube_config_temp_dir_path(arguments: &Arguments) -> Result<String> {
     Ok(unique_kube_config_temp_dir)
 }
 
+fn args_validation(args: &Arguments) -> Result<()> {
+    match &args.subcommand {
+        SubCommand::IntegrationTest(integ_test_args) => {
+            ensure!(
+                ARCHES.contains(&ArchitectureValues::from(
+                    integ_test_args.ami_arch.as_str().clone()
+                )),
+                error::InvalidArchInput {
+                    input: integ_test_args.ami_arch.clone()
+                }
+            )
+        }
+        _ => return Ok(()),
+    }
+    Ok(())
+}
+
 async fn run() -> Result<()> {
     // Parse and store the args passed to the program
     let args = Arguments::from_args();
 
+    // Validate the args
+    args_validation(&args)?;
+
     let cluster_info = get_cluster_info(&args.cluster_name, &args.region)
         .await
         .context(error::GetClusterInfo)?;
@@ -122,10 +156,13 @@ async fn run() -> Result<()> {
         SubCommand::IntegrationTest(integ_test_args) => {
             // create instances and add nodes to eks cluster
             info!("Creating EC2 instances ...");
-            let created_instances =
-                create_ec2_instance(cluster_info, &integ_test_args.instance_ami_id)
-                    .await
-                    .context(error::CreateEc2Instances)?;
+            let created_instances = create_ec2_instance(
+                cluster_info,
+                &integ_test_args.ami_arch,
+                &integ_test_args.bottlerocket_version,
+            )
+            .await
+            .context(error::CreateEc2Instances)?;
             info!("EC2 instances have been created");
 
             // generate kubeconfig if no input value for argument `kube_config_path`
@@ -201,6 +238,9 @@ mod error {
         #[snafu(display("Failed to create ec2 instances: {}", source))]
         CreateEc2Instances { source: ProviderError },
 
+        #[snafu(display("Invalid Arch input: {}", input))]
+        InvalidArchInput { input: String },
+
         #[snafu(display("Failed to label ec2 instances: {}", source))]
         LabelNode { source: update_error::Error },