diff --git a/src/bin/mock.rs b/src/bin/mock.rs
index eecf1ac..9bbcdf9 100644
--- a/src/bin/mock.rs
+++ b/src/bin/mock.rs
@@ -345,6 +345,7 @@ mod tests {
     // use scheduler2::parser::ExecutionPlanParser;
     // use crate::{file_mode, run_single_query, start_system, CATALOG_PATH, TPCH_FILES};
     use crate::TPCH_FILES;
+    use crate::file_mode;
 
     #[tokio::test]
     async fn test_file_mode() {
diff --git a/src/integration_test.rs b/src/integration_test.rs
index cd46762..55c7dfd 100644
--- a/src/integration_test.rs
+++ b/src/integration_test.rs
@@ -23,7 +23,7 @@ pub struct IntegrationTest {
 }
 
 const CONFIG_PATH: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/executors.toml");
-const CATALOG_PATH: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/test_data");
+pub const CATALOG_PATH: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/test_data");
 const LOG_PATH: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/executor_logs");
 
 impl IntegrationTest {
@@ -195,7 +195,8 @@ impl IntegrationTest {
 mod tests {
     use crate::integration_test::IntegrationTest;
     use crate::parser::ExecutionPlanParser;
-    use crate::CATALOG_PATH;
+    // use crate::CATALOG_PATH;
+    use super::*;
     use datafusion::arrow::array::{Int32Array, RecordBatch};
     use datafusion::arrow::datatypes::{DataType, Field, Schema};
     use std::path::PathBuf;
diff --git a/src/intermediate_results.rs b/src/intermediate_results.rs
index 9e817c0..dfd979a 100644
--- a/src/intermediate_results.rs
+++ b/src/intermediate_results.rs
@@ -112,8 +112,8 @@ pub async fn rewrite_query(
 mod tests {
     use super::*;
     use crate::parser::ExecutionPlanParser;
-    use crate::CATALOG_PATH;
-    use datafusion::arrow::array::{Array, Int32Array};
+    use crate::integration_test::CATALOG_PATH;
+    use datafusion::arrow::array::Int32Array;
     use datafusion::arrow::datatypes::{DataType, Field, Schema};
     use datafusion::arrow::record_batch::RecordBatch;
 
diff --git a/src/queue.rs b/src/queue.rs
index c216a6e..983581f 100644
--- a/src/queue.rs
+++ b/src/queue.rs
@@ -187,6 +187,8 @@ impl Queue {
 mod tests {
     use rand::Rng;
     use tokio::sync::{Mutex, Notify};
+    use tokio::time::sleep;
+    use std::time::Duration;
 
     use crate::parser::ExecutionPlanParser;
     use crate::{
@@ -203,18 +205,30 @@ mod tests {
         time::SystemTime,
     };
 
+    // Test that query keys compare properly.
     #[tokio::test]
     async fn test_query_key_cmp() {
         let then = SystemTime::now();
+        let now1 = SystemTime::now().duration_since(then).unwrap();
+        sleep(Duration::from_secs(1)).await;
+        let now2 = SystemTime::now().duration_since(then).unwrap();
+
         let key1 = QueryKey {
-            ft: SystemTime::now().duration_since(then).unwrap(),
-            qid: 0,
+            ft: now1.clone(),
+            qid: 1,
         };
         let key2 = QueryKey {
-            ft: SystemTime::now().duration_since(then).unwrap(),
+            ft: now2,
+            qid: 0,
+        };
+        let key3 = QueryKey {
+            ft: now1,
             qid: 0,
         };
+        // Make sure durations are compared first
         assert!(key1 < key2);
+        // Then qids
+        assert!(key3 < key1);
     }
 
     #[tokio::test]
diff --git a/src/task_queue.rs b/src/task_queue.rs
index 2e419a5..9c2513b 100644
--- a/src/task_queue.rs
+++ b/src/task_queue.rs
@@ -35,7 +35,7 @@ impl TaskQueue {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::server::composable_database::TaskId;
+    use crate::composable_database::TaskId;
     use crate::task::TaskStatus;
     use std::sync::Arc;
     use tokio::runtime::Runtime;