Update onefuzz-agent unit tests (#592)

src/agent/onefuzz-agent/src/tasks/fuzz/supervisor.rs

@@ -226,7 +226,7 @@ mod tests {
     use super::*;
     use crate::tasks::stats::afl::read_stats;
     use onefuzz::process::monitor_process;
-    use onefuzz::telemetry::EventData;
+    use onefuzz_telemetry::EventData;
     use std::collections::HashMap;
     use std::time::Instant;
 

src/agent/onefuzz-agent/src/tasks/generic/input_poller.rs

@@ -27,6 +27,23 @@ pub enum State<M> {
     Processed(M),
 }
 
+impl<M: PartialEq> PartialEq for State<M> {
+    fn eq(&self, other: &State<M>) -> bool {
+        use State::*;
+
+        match (self, other) {
+            (Ready, Ready) => true,
+            (Polled(l), Polled(r)) => l == r,
+            (Parsed(l0, l1), Parsed(r0, r1)) => l0 == r0 && l1 == r1,
+            (Downloaded(l0, l1, l2, l3), Downloaded(r0, r1, r2, r3)) => {
+                l0 == r0 && l1 == r1 && l2 == r2 && l3.path() == r3.path()
+            }
+            (Processed(l), Processed(r)) => l == r,
+            _ => false,
+        }
+    }
+}
+
 impl<M> fmt::Display for State<M> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {

src/agent/onefuzz-agent/src/tasks/generic/input_poller/tests.rs

@@ -180,7 +180,7 @@ async fn test_parsed_download() {
         .unwrap();
 
     match task.state() {
-        State::Downloaded(got_msg, got_url, got_path) => {
+        State::Downloaded(got_msg, got_url, got_path, _tmp_dir) => {
             assert_eq!(*got_msg, msg);
             assert_eq!(*got_url, url);
             assert_eq!(got_path.file_name(), input.file_name());
@@ -194,6 +194,7 @@ async fn test_parsed_download() {
 #[tokio::test]
 async fn test_downloaded_process() {
     let mut task = fixture();
+    let tmp_dir = tempfile::tempdir().unwrap();
 
     let dir = Path::new("etc");
 
@@ -201,7 +202,7 @@ async fn test_downloaded_process() {
     let url = url_fixture(msg);
     let input = input_fixture(dir, msg);
 
-    task.set_state(State::Downloaded(msg, url.clone(), input.clone()));
+    task.set_state(State::Downloaded(msg, url.clone(), input.clone(), tmp_dir));
 
     let mut processor = TestProcessor::default();