fix unit tests

Author: noot

crates/worker/src/cli/command.rs

@@ -220,11 +220,19 @@ pub async fn execute_command(
                 );
                 std::process::exit(1);
             }
-            let state = Arc::new(SystemState::new(
+            let state = match SystemState::new(
                 state_dir_overwrite.clone(),
                 *disable_state_storing,
                 *compute_pool_id,
-            ));
+            ) {
+                Ok(state) => state,
+                Err(e) => {
+                    error!("❌ Failed to initialize system state: {e}");
+                    std::process::exit(1);
+                }
+            };
+
+            let state = Arc::new(state);
 
             let private_key_provider = if let Some(key) = private_key_provider {
                 Console::warning("Using private key from command line is not recommended. Consider using PRIVATE_KEY_PROVIDER environment variable instead.");

crates/worker/src/docker/taskbridge/bridge.rs

@@ -473,7 +473,7 @@ mod tests {
         let temp_dir = tempdir()?;
         let socket_path = temp_dir.path().join("test.sock");
         let metrics_store = Arc::new(MetricsStore::new());
-        let state = Arc::new(SystemState::new(None, false, 0));
+        let state = Arc::new(SystemState::new(None, false, 0).unwrap());
         let bridge = TaskBridge::new(
             Some(socket_path.to_str().unwrap()),
             metrics_store.clone(),
@@ -506,7 +506,7 @@ mod tests {
         let temp_dir = tempdir()?;
         let socket_path = temp_dir.path().join("test.sock");
         let metrics_store = Arc::new(MetricsStore::new());
-        let state = Arc::new(SystemState::new(None, false, 0));
+        let state = Arc::new(SystemState::new(None, false, 0).unwrap());
         let bridge = TaskBridge::new(
             Some(socket_path.to_str().unwrap()),
             metrics_store.clone(),
@@ -541,7 +541,7 @@ mod tests {
         let temp_dir = tempdir()?;
         let socket_path = temp_dir.path().join("test.sock");
         let metrics_store = Arc::new(MetricsStore::new());
-        let state = Arc::new(SystemState::new(None, false, 0));
+        let state = Arc::new(SystemState::new(None, false, 0).unwrap());
         let bridge = TaskBridge::new(
             Some(socket_path.to_str().unwrap()),
             metrics_store.clone(),
@@ -590,7 +590,7 @@ mod tests {
         let temp_dir = tempdir()?;
         let socket_path = temp_dir.path().join("test.sock");
         let metrics_store = Arc::new(MetricsStore::new());
-        let state = Arc::new(SystemState::new(None, false, 0));
+        let state = Arc::new(SystemState::new(None, false, 0).unwrap());
         let bridge = TaskBridge::new(
             Some(socket_path.to_str().unwrap()),
             metrics_store.clone(),
@@ -639,7 +639,7 @@ mod tests {
         let temp_dir = tempdir()?;
         let socket_path = temp_dir.path().join("test.sock");
         let metrics_store = Arc::new(MetricsStore::new());
-        let state = Arc::new(SystemState::new(None, false, 0));
+        let state = Arc::new(SystemState::new(None, false, 0).unwrap());
         let bridge = TaskBridge::new(
             Some(socket_path.to_str().unwrap()),
             metrics_store.clone(),

crates/worker/src/state/system_state.rs

@@ -1,3 +1,4 @@
+use anyhow::bail;
 use anyhow::Result;
 use directories::ProjectDirs;
 use log::debug;
@@ -60,7 +61,7 @@ impl SystemState {
         state_dir: Option<String>,
         disable_state_storing: bool,
         compute_pool_id: u32,
-    ) -> Self {
+    ) -> Result<Self> {
         let default_state_dir = get_default_state_dir();
         debug!("Default state dir: {default_state_dir:?}");
         let state_path = state_dir
@@ -84,7 +85,7 @@ impl SystemState {
                     endpoint = loaded_state.endpoint;
                     p2p_keypair = Some(loaded_state.p2p_keypair);
                 } else {
-                    debug!("Failed to load state from {state_file:?}");
+                    bail!("failed to load state from {state_file:?}");
                 }
             }
         }
@@ -93,15 +94,15 @@ impl SystemState {
             p2p_keypair = Some(p2p::Keypair::generate_ed25519());
         }
 
-        Self {
+        Ok(Self {
             last_heartbeat: Arc::new(RwLock::new(None)),
             is_running: Arc::new(RwLock::new(false)),
             endpoint: Arc::new(RwLock::new(endpoint)),
             state_dir_overwrite: state_path.clone(),
             disable_state_storing,
             compute_pool_id,
             p2p_keypair: p2p_keypair.expect("p2p keypair must be Some at this point"),
-        }
+        })
     }
 
     fn save_state(&self, heartbeat_endpoint: Option<String>) -> Result<()> {
@@ -141,8 +142,7 @@ impl SystemState {
             match serde_json::from_str(&contents) {
                 Ok(state) => return Ok(Some(state)),
                 Err(e) => {
-                    debug!("Error parsing state file: {e}");
-                    return Ok(None);
+                    bail!("failed to parse state file: {e}");
                 }
             }
         }
@@ -232,7 +232,8 @@ mod tests {
             Some(temp_dir.path().to_string_lossy().to_string()),
             false,
             0,
-        );
+        )
+        .unwrap();
         let _ = state
             .set_running(true, Some("http://localhost:8080/heartbeat".to_string()))
             .await;
@@ -255,30 +256,33 @@ mod tests {
         let state_file = temp_dir.path().join(STATE_FILENAME);
         fs::write(&state_file, "invalid_toml_content").expect("Failed to write to state file");
 
-        let state = SystemState::new(
+        assert!(SystemState::new(
             Some(temp_dir.path().to_string_lossy().to_string()),
             false,
             0,
-        );
-        assert!(!(state.is_running().await));
-        assert_eq!(state.get_heartbeat_endpoint().await, None);
+        )
+        .is_err());
     }
 
     #[tokio::test]
     async fn test_load_state() {
+        let keypair = p2p::Keypair::generate_ed25519();
+        let state = PersistedSystemState {
+            endpoint: Some("http://localhost:8080/heartbeat".to_string()),
+            p2p_keypair: keypair,
+        };
+        let serialized = serde_json::to_string_pretty(&state).unwrap();
+
         let temp_dir = setup_test_dir();
         let state_file = temp_dir.path().join(STATE_FILENAME);
-        fs::write(
-            &state_file,
-            r#"{"endpoint": "http://localhost:8080/heartbeat"}"#,
-        )
-        .expect("Failed to write to state file");
+        fs::write(&state_file, serialized).unwrap();
 
         let state = SystemState::new(
             Some(temp_dir.path().to_string_lossy().to_string()),
             false,
             0,
-        );
+        )
+        .unwrap();
         assert_eq!(
             state.get_heartbeat_endpoint().await,
             Some("http://localhost:8080/heartbeat".to_string())