Fix rest handler update test

Author: rdettai

quickwit/Cargo.lock

@@ -47,6 +47,7 @@ tracing = { workspace = true }
 ulid = { workspace = true }
 utoipa = { workspace = true }
 vrl = { workspace = true, optional = true }
+warp = { workspace = true, optional = true }
 
 quickwit-actors = { workspace = true }
 quickwit-aws = { workspace = true }
@@ -84,7 +85,8 @@ sqs = [
   "queue-sources",
   "quickwit-aws/sqs",
 ]
-sqs-localstack-tests = []
+sqs-test-helpers = ["warp"]
+sqs-localstack-tests = ["sqs-test-helpers"]
 vendored-kafka = [
   "kafka",
   "libz-sys/static",

quickwit/quickwit-indexing/Cargo.toml

@@ -261,10 +261,11 @@ pub(crate) async fn check_connectivity(queue_url: &str) -> anyhow::Result<()> {
     Ok(())
 }
 
-#[cfg(feature = "sqs-localstack-tests")]
+#[cfg(feature = "sqs-test-helpers")]
 pub mod test_helpers {
     use aws_sdk_sqs::types::QueueAttributeName;
     use ulid::Ulid;
+    use warp::Filter;
 
     use super::*;
 
@@ -316,6 +317,31 @@ pub mod test_helpers {
             .unwrap()
             .to_string()
     }
+
+    /// Runs a mock SQS GetQueueAttributes endpoint to enable creating SQS
+    /// sources that pass the connectivity check
+    ///
+    /// Rerturns the queue URL to use for the source and a guard for the
+    /// temporary mock server
+    pub fn start_mock_sqs_get_queue_attributes_endpoint() -> (String, oneshot::Sender<()>) {
+        let hello = warp::path!().map(|| "{}");
+        let (tx, rx) = oneshot::channel();
+        let (addr, server) =
+            warp::serve(hello).bind_with_graceful_shutdown(([127, 0, 0, 1], 0), async {
+                rx.await.ok();
+            });
+        tokio::spawn(server);
+        let queue_url = format!("http://{}:{}/", addr.ip(), addr.port());
+        (queue_url, tx)
+    }
+
+    #[tokio::test]
+    async fn test_mock_sqs_get_queue_attributes_endpoint() {
+        let (queue_url, _shutdown) = start_mock_sqs_get_queue_attributes_endpoint();
+        check_connectivity(&queue_url).await.unwrap();
+        drop(_shutdown);
+        check_connectivity(&queue_url).await.unwrap_err();
+    }
 }
 
 #[cfg(test)]

quickwit/quickwit-indexing/src/source/queue_sources/sqs_queue.rs

@@ -98,3 +98,7 @@ pprof = [
   "dep:pprof"
 ]
 testsuite = []
+sqs-for-tests = [
+  "quickwit-indexing/sqs",
+  "quickwit-indexing/sqs-test-helpers"
+]

quickwit/quickwit-serve/Cargo.toml

@@ -1111,9 +1111,13 @@ mod tests {
         }
     }
 
+    #[cfg(feature = "sqs-for-tests")]
     #[tokio::test]
     async fn test_update_source() {
+        use quickwit_indexing::source::sqs_queue::test_helpers::start_mock_sqs_get_queue_attributes_endpoint;
+
         let metastore = metastore_for_test();
+        let (queue_url, _guard) = start_mock_sqs_get_queue_attributes_endpoint();
         let index_service = IndexService::new(metastore.clone(), StorageResolver::unconfigured());
         let mut node_config = NodeConfig::for_test();
         node_config.default_index_root_uri = Uri::for_test("file:///default-index-root-uri");
@@ -1137,64 +1141,77 @@ mod tests {
         assert_json_include!(actual: resp_json, expected: expected_response_json);
 
         // Create source.
-        let source_config_body = r#"{"version": "0.7", "source_id": "vec-source", "source_type": "vec", "params": {"docs": [], "batch_num_docs": 10}}"#;
+        let source_config_body = serde_json::json!({
+            "version": "0.7",
+            "source_id": "sqs-source",
+            "source_type": "file",
+            "params": {"notifications": [{"type": "sqs", "queue_url": queue_url, "message_type": "s3_notification"}]},
+        });
         let resp = warp::test::request()
             .path("/indexes/hdfs-logs/sources")
             .method("POST")
-            .json(&true)
-            .body(source_config_body)
+            .json(&source_config_body)
             .reply(&index_management_handler)
             .await;
-        assert_eq!(resp.status(), 200);
-
-        // Update the source.
-        let update_source_config_body = r#"{"version": "0.7", "source_id": "vec-source", "source_type": "vec", "params": {"docs": [], "batch_num_docs": 20}}"#;
-        let resp = warp::test::request()
-            .path("/indexes/hdfs-logs/sources/vec-source")
-            .method("PUT")
-            .json(&true)
-            .body(update_source_config_body)
-            .reply(&index_management_handler)
-            .await;
-        assert_eq!(resp.status(), 200);
-        // Check that the source has been updated.
-        let index_metadata = metastore
-            .index_metadata(IndexMetadataRequest::for_index_id("hdfs-logs".to_string()))
-            .await
-            .unwrap()
-            .deserialize_index_metadata()
-            .unwrap();
-        assert!(index_metadata.sources.contains_key("vec-source"));
-        let source_config = index_metadata.sources.get("vec-source").unwrap();
-        assert_eq!(source_config.source_type(), SourceType::Vec);
-        assert_eq!(
-            source_config.source_params,
-            SourceParams::Vec(VecSourceParams {
-                docs: Vec::new(),
-                batch_num_docs: 20,
-                partition: "".to_string(),
-            })
-        );
+        let resp_body = std::str::from_utf8(resp.body()).unwrap();
+        assert_eq!(resp.status(), 200, "{resp_body}");
 
-        // Update the source with a different source_id (forbidden)
-        let update_source_config_body = r#"{"version": "0.7", "source_id": "other-source-id", "source_type": "vec", "params": {"docs": [], "batch_num_docs": 20}}"#;
-        let resp = warp::test::request()
-            .path("/indexes/hdfs-logs/sources/vec-source")
-            .method("PUT")
-            .json(&true)
-            .body(update_source_config_body)
-            .reply(&index_management_handler)
-            .await;
-        assert_eq!(resp.status(), 400);
-        // Check that the source hasn't been updated.
-        let index_metadata = metastore
-            .index_metadata(IndexMetadataRequest::for_index_id("hdfs-logs".to_string()))
-            .await
-            .unwrap()
-            .deserialize_index_metadata()
-            .unwrap();
-        assert!(index_metadata.sources.contains_key("vec-source"));
-        assert!(!index_metadata.sources.contains_key("other-source-id"));
+        {
+            // Update the source.
+            let update_source_config_body = serde_json::json!({
+                "version": "0.7",
+                "source_id": "sqs-source",
+                "source_type": "file",
+                "params": {"notifications": [{"type": "sqs", "queue_url": queue_url, "message_type": "s3_notification"}]},
+            });
+            let resp = warp::test::request()
+                .path("/indexes/hdfs-logs/sources/sqs-source")
+                .method("PUT")
+                .json(&update_source_config_body)
+                .reply(&index_management_handler)
+                .await;
+            let resp_body = std::str::from_utf8(resp.body()).unwrap();
+            assert_eq!(resp.status(), 200, "{resp_body}");
+            // Check that the source has been updated.
+            let index_metadata = metastore
+                .index_metadata(IndexMetadataRequest::for_index_id("hdfs-logs".to_string()))
+                .await
+                .unwrap()
+                .deserialize_index_metadata()
+                .unwrap();
+            let metastore_source_config = index_metadata.sources.get("sqs-source").unwrap();
+            assert_eq!(metastore_source_config.source_type(), SourceType::File);
+            assert_eq!(
+                metastore_source_config,
+                &serde_json::from_value(update_source_config_body).unwrap(),
+            );
+        }
+        {
+            // Update the source with a different source_id (forbidden)
+            let update_source_config_body = serde_json::json!({
+                "version": "0.7",
+                "source_id": "new-source-id",
+                "source_type": "file",
+                "params": {"notifications": [{"type": "sqs", "queue_url": queue_url, "message_type": "s3_notification"}]},
+            });
+            let resp = warp::test::request()
+                .path("/indexes/hdfs-logs/sources/sqs-source")
+                .method("PUT")
+                .json(&update_source_config_body)
+                .reply(&index_management_handler)
+                .await;
+            let resp_body = std::str::from_utf8(resp.body()).unwrap();
+            assert_eq!(resp.status(), 400, "{resp_body}");
+            // Check that the source hasn't been updated.
+            let index_metadata = metastore
+                .index_metadata(IndexMetadataRequest::for_index_id("hdfs-logs".to_string()))
+                .await
+                .unwrap()
+                .deserialize_index_metadata()
+                .unwrap();
+            assert!(index_metadata.sources.contains_key("sqs-source"));
+            assert!(!index_metadata.sources.contains_key("other-source-id"));
+        }
     }
 
     #[tokio::test]