From aa7415f90d9d385822f26f18581a19a68125d1f0 Mon Sep 17 00:00:00 2001
From: Kanan Boubion <kananboubion@microsoft.com>
Date: Wed, 4 Oct 2023 15:22:48 -0700
Subject: [PATCH] Fix import issues

---
 src/agent/onefuzz-task/src/tasks/config.rs | 33 ++++++----------------
 src/agent/onefuzz-task/src/test_utils.rs   |  1 -
 2 files changed, 8 insertions(+), 26 deletions(-)

diff --git a/src/agent/onefuzz-task/src/tasks/config.rs b/src/agent/onefuzz-task/src/tasks/config.rs
index 2c4e76310..4ee71d935 100644
--- a/src/agent/onefuzz-task/src/tasks/config.rs
+++ b/src/agent/onefuzz-task/src/tasks/config.rs
@@ -395,7 +395,6 @@ mod tests {
     use proptest::prelude::*;
     use onefuzz::expand::GetExpand;
     use std::collections::HashMap;
-    use crate::test_utils::arbitraries::*;
 
     use super::CommonConfig;
 
@@ -416,33 +415,17 @@ mod tests {
                 ("task_id", config.task_id.to_string()),
                 ("setup_dir", dunce::canonicalize(config.setup_dir.clone()).unwrap().to_string_lossy().to_string()),
             ]);
-            // Look for a shorthand way of doing these optional hashmap inserts
-            // config.instance_telemetry_key.iter().for_each(|key| {
-            //     params.insert("instance_telemetry_key", key.to_string());
-            // });
-            match &config.instance_telemetry_key {
-                Some(key) => {
-                    params.insert("instance_telemetry_key", key.clone().to_string());
-                },
-                None => {},
+            if let Some(key) = &config.instance_telemetry_key {
+                params.insert("instance_telemetry_key", key.to_string());
             }
-            match &config.microsoft_telemetry_key {
-                Some(key) => {
-                    params.insert("microsoft_telemetry_key", key.clone().to_string());
-                },
-                None => {},
+            if let Some(key) = &config.microsoft_telemetry_key {
+                params.insert("microsoft_telemetry_key", key.clone().to_string());
             }
-            match &config.extra_setup_dir {
-                Some(dir) => {
-                    params.insert("extra_setup_dir", dunce::canonicalize(dir).unwrap().to_string_lossy().to_string());
-                },
-                None => {},
+            if let Some(dir) = &config.extra_setup_dir {
+                params.insert("extra_setup_dir", dunce::canonicalize(dir).unwrap().to_string_lossy().to_string());
             }
-            match &config.extra_output {
-                Some(dir) => {
-                    params.insert("extra_output_dir", dunce::canonicalize(&dir.local_path).unwrap().to_string_lossy().to_string());
-                },
-                None => {},
+            if let Some(dir) = &config.extra_output {
+                params.insert("extra_output_dir", dunce::canonicalize(&dir.local_path).unwrap().to_string_lossy().to_string());
             }
 
             for (param, expected) in params.iter() {
diff --git a/src/agent/onefuzz-task/src/test_utils.rs b/src/agent/onefuzz-task/src/test_utils.rs
index d8ad0fbe5..29f8ce16d 100644
--- a/src/agent/onefuzz-task/src/test_utils.rs
+++ b/src/agent/onefuzz-task/src/test_utils.rs
@@ -1,4 +1,3 @@
-#[cfg(test)]
 pub mod arbitraries {
     use std::path::PathBuf;