From 0c014828a7f7dc2fdbaf1fb956b2a5c3ae429e27 Mon Sep 17 00:00:00 2001
From: Martin Moreira de Jesus <martin.moreira-de-jesus@protonmail.com>
Date: Tue, 30 Apr 2024 17:56:26 +0200
Subject: [PATCH] fix(agent): remove unwraps

---
 src/agent/src/agents/mod.rs       | 11 +++++++----
 src/agent/src/lib.rs              |  2 ++
 src/agent/src/workload/config.rs  |  8 ++++----
 src/agent/src/workload/runner.rs  |  8 ++++----
 src/agent/src/workload/service.rs |  3 ++-
 5 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/src/agent/src/agents/mod.rs b/src/agent/src/agents/mod.rs
index 13e3fff..f70d43b 100644
--- a/src/agent/src/agents/mod.rs
+++ b/src/agent/src/agents/mod.rs
@@ -1,4 +1,4 @@
-use crate::AgentResult;
+use crate::{AgentError, AgentResult};
 use serde::Deserialize;
 
 #[cfg(feature = "debug-agent")]
@@ -36,14 +36,17 @@ impl std::fmt::Display for Language {
 }
 
 impl TryFrom<&str> for Language {
-    type Error = String;
+    type Error = AgentError;
 
-    fn try_from(value: &str) -> Result<Self, Self::Error> {
+    fn try_from(value: &str) -> Result<Self, AgentError> {
         match value {
             "rust" => Ok(Language::Rust),
             #[cfg(feature = "debug-agent")]
             "debug" => Ok(Language::Debug),
-            _ => Err(format!("Invalid language: {}", value)),
+            _ => Err(AgentError::InvalidLanguage(format!(
+                "Invalid language: {}",
+                value
+            ))),
         }
     }
 }
diff --git a/src/agent/src/lib.rs b/src/agent/src/lib.rs
index 9332611..0c43daf 100644
--- a/src/agent/src/lib.rs
+++ b/src/agent/src/lib.rs
@@ -8,6 +8,7 @@ pub mod workload;
 pub enum AgentError {
     OpenConfigFileError(std::io::Error),
     ParseConfigError(toml::de::Error),
+    InvalidLanguage(String),
     BuildFailed(AgentOutput),
 }
 
@@ -17,6 +18,7 @@ impl fmt::Display for AgentError {
             AgentError::OpenConfigFileError(e) => write!(f, "Failed to open config file: {}", e),
             AgentError::ParseConfigError(e) => write!(f, "Failed to parse config file: {}", e),
             AgentError::BuildFailed(output) => write!(f, "Build failed: {:?}", output),
+            AgentError::InvalidLanguage(e) => write!(f, "Invalid language: {}", e),
         }
     }
 }
diff --git a/src/agent/src/workload/config.rs b/src/agent/src/workload/config.rs
index 40261b4..f93b336 100644
--- a/src/agent/src/workload/config.rs
+++ b/src/agent/src/workload/config.rs
@@ -35,15 +35,15 @@ impl Config {
         Ok(config)
     }
 
-    pub fn new_from_execute_request(execute_request: ExecuteRequest) -> Self {
-        Self {
+    pub fn new_from_execute_request(execute_request: ExecuteRequest) -> Result<Self, AgentError> {
+        Ok(Self {
             workload_name: execute_request.workload_name.clone(),
             // TODO: Fix this unwrap
-            language: Language::try_from(execute_request.language.clone().as_str()).unwrap(),
+            language: Language::try_from(execute_request.language.clone().as_str())?,
             action: execute_request.action().into(),
             config_string: execute_request.config_str,
             code: execute_request.code,
-        }
+        })
     }
 }
 
diff --git a/src/agent/src/workload/runner.rs b/src/agent/src/workload/runner.rs
index 47f4ba2..ae4e175 100644
--- a/src/agent/src/workload/runner.rs
+++ b/src/agent/src/workload/runner.rs
@@ -2,7 +2,7 @@ use crate::{
     agent::ExecuteRequest,
     agents::{rust, Agent, AgentOutput, Language},
     workload::config::Action,
-    AgentResult,
+    AgentError, AgentResult,
 };
 
 #[cfg(feature = "debug-agent")]
@@ -28,9 +28,9 @@ impl Runner {
         Runner { config, agent }
     }
 
-    pub fn new_from_execute_request(execute_request: ExecuteRequest) -> Self {
-        let config = Config::new_from_execute_request(execute_request);
-        Self::new(config)
+    pub fn new_from_execute_request(execute_request: ExecuteRequest) -> Result<Self, AgentError> {
+        let config = Config::new_from_execute_request(execute_request)?;
+        Ok(Self::new(config))
     }
 
     pub fn run(&self) -> AgentResult<AgentOutput> {
diff --git a/src/agent/src/workload/service.rs b/src/agent/src/workload/service.rs
index 90ae1ff..3b05386 100644
--- a/src/agent/src/workload/service.rs
+++ b/src/agent/src/workload/service.rs
@@ -18,7 +18,8 @@ impl WorkloadRunner for WorkloadRunnerService {
 
         let execute_request = req.into_inner();
 
-        let runner = Runner::new_from_execute_request(execute_request);
+        let runner = Runner::new_from_execute_request(execute_request)
+            .map_err(|e| tonic::Status::internal(e.to_string()))?;
 
         let res = runner
             .run()