From 30f3b05b42b35a10b87cf412a78c5fd6a4c1f52c Mon Sep 17 00:00:00 2001
From: Erdem Meydanli <meydanli@amazon.com>
Date: Tue, 19 Mar 2024 22:06:02 +0000
Subject: [PATCH] docker.rs: add more tests

Add simple tests to validate functionality.

Signed-off-by: Erdem Meydanli <meydanli@amazon.com>
---
 enclave_build/src/docker.rs | 66 +++++++++++++++++++++++++++++++++++--
 1 file changed, 64 insertions(+), 2 deletions(-)

diff --git a/enclave_build/src/docker.rs b/enclave_build/src/docker.rs
index 297afd32d..348d7123e 100644
--- a/enclave_build/src/docker.rs
+++ b/enclave_build/src/docker.rs
@@ -430,14 +430,15 @@ fn write_config(config: Vec<String>) -> Result<NamedTempFile, DockerError> {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use std::io::Read;
+    use base64::engine::general_purpose;
+    use std::{env, io::Read};
 
     /// Test extracted configuration is as expected
     #[test]
     fn test_config() {
         let docker = DockerUtil::new(String::from("public.ecr.aws/aws-nitro-enclaves/hello:v1"));
 
-        let (cmd_file, env_file) = docker.load().unwrap();
+        let (cmd_file, env_file) = docker.unwrap().load().unwrap();
         let mut cmd_file = File::open(cmd_file.path()).unwrap();
         let mut env_file = File::open(env_file.path()).unwrap();
 
@@ -453,4 +454,65 @@ mod tests {
              HELLO=Hello from the enclave side!\n"
         );
     }
+
+    #[test]
+    fn test_new() {
+        let docker = DockerUtil::new(String::from("alpine")).unwrap();
+        assert_eq!(docker.docker_image, "alpine:latest");
+        let docker = DockerUtil::new(String::from("nginx:1.19")).unwrap();
+        assert_eq!(docker.docker_image, "nginx:1.19");
+    }
+
+    #[test]
+    fn test_get_credentials() {
+        let test_user = "test_user";
+        let test_password = "test_password";
+        let auth = format!("{}:{}", test_user, test_password);
+        let encoded_auth = general_purpose::STANDARD.encode(auth);
+        let config = format!(
+            r#"{{
+            "auths": {{
+              "https://public.ecr.aws/aws-nitro-enclaves/hello/v1/": {{
+                "auth": "{}"
+              }},
+              "https://registry.example.com": {{
+                "auth": "b3RoZXJfdXNlcjpvdGhlcl9wYXNzd29yZA=="
+              }}
+            }}
+          }}"#,
+            encoded_auth
+        );
+
+        // Create a temporary file
+        let mut temp_file = NamedTempFile::new().expect("Failed to create temporary file.");
+
+        // Write the config to the temporary file
+        write!(temp_file, "{}", config).expect("Failed to write to temporary file.");
+
+        // Set the DOCKER_CONFIG environment variable to point to the temporary file's path
+        let temp_file_path = temp_file.path().to_string_lossy().to_string();
+        env::set_var("DOCKER_CONFIG", temp_file_path);
+
+        let docker =
+            DockerUtil::new(String::from("public.ecr.aws/aws-nitro-enclaves/hello:v1")).unwrap();
+        let creds = docker.get_credentials().unwrap();
+        assert_eq!(creds.username, Some(test_user.to_string()));
+        assert_eq!(creds.password, Some(test_password.to_string()));
+
+        temp_file.close().unwrap();
+    }
+
+    #[test]
+    fn test_architecture() {
+        let docker =
+            DockerUtil::new(String::from("public.ecr.aws/aws-nitro-enclaves/hello:v1")).unwrap();
+        docker.pull_image().unwrap();
+        let arch = docker.architecture().unwrap();
+        assert_eq!(arch, "amd64");
+
+        let docker = DockerUtil::new(String::from("arm64v8/alpine")).unwrap();
+        docker.pull_image().unwrap();
+        let arch = docker.architecture().unwrap();
+        assert_eq!(arch, "arm64");
+    }
 }