diff --git a/src/agent/onefuzz-agent/src/tasks/fuzz/supervisor.rs b/src/agent/onefuzz-agent/src/tasks/fuzz/supervisor.rs
index 0ba1393c3..81ed014f7 100644
--- a/src/agent/onefuzz-agent/src/tasks/fuzz/supervisor.rs
+++ b/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;
 
diff --git a/src/agent/onefuzz-agent/src/tasks/generic/input_poller.rs b/src/agent/onefuzz-agent/src/tasks/generic/input_poller.rs
index 94eedf8e0..3361cce29 100644
--- a/src/agent/onefuzz-agent/src/tasks/generic/input_poller.rs
+++ b/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 {
diff --git a/src/agent/onefuzz-agent/src/tasks/generic/input_poller/tests.rs b/src/agent/onefuzz-agent/src/tasks/generic/input_poller/tests.rs
index 464e64432..c4a83422b 100644
--- a/src/agent/onefuzz-agent/src/tasks/generic/input_poller/tests.rs
+++ b/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();
 
diff --git a/src/agent/reqwest-retry/src/lib.rs b/src/agent/reqwest-retry/src/lib.rs
index 092b1af6a..b65005a92 100644
--- a/src/agent/reqwest-retry/src/lib.rs
+++ b/src/agent/reqwest-retry/src/lib.rs
@@ -173,6 +173,8 @@ impl SendRetry for reqwest::RequestBuilder {
 mod test {
     use super::*;
 
+    // TODO: convert to feature-gated integration test.
+    #[ignore]
     #[tokio::test]
     async fn empty_stack() -> Result<()> {
         let resp = reqwest::Client::new()
diff --git a/src/ci/agent.sh b/src/ci/agent.sh
index 5abcb85e4..0577a21ef 100755
--- a/src/ci/agent.sh
+++ b/src/ci/agent.sh
@@ -15,7 +15,7 @@ if sccache --help; then
 fi
 
 # only set CARGO_INCREMENTAL on non-release builds
-# 
+#
 # This speeds up build time, but makes the resulting binaries slightly slower.
 # https://doc.rust-lang.org/cargo/reference/profiles.html?highlight=incremental#incremental
 if [ "${GITHUB_REF}" != "" ]; then
@@ -38,8 +38,11 @@ cargo clippy --release -- -D warnings
 # export RUST_LOG=trace
 export RUST_BACKTRACE=full
 cargo test --release --manifest-path ./onefuzz-supervisor/Cargo.toml
-# TODO: re-enable
+
+# TODO: re-enable integration tests.
 # cargo test --release --manifest-path ./onefuzz-agent/Cargo.toml --features integration_test -- --nocapture
+cargo test --release --manifest-path ./onefuzz-agent/Cargo.toml
+
 cargo test --release --manifest-path ./onefuzz/Cargo.toml
 
 cp target/release/onefuzz-agent* ../../artifacts/agent