Temporary fix

timonv · timonv · commit f7559bbfd301 · 2025-02-13T14:33:02.000+01:00
diff --git a/swiftide-integrations/src/scraping/loader.rs b/swiftide-integrations/src/scraping/loader.rs
@@ -1,8 +1,5 @@
-use std::sync::Arc;
-
 use derive_builder::Builder;
 use spider::website::Website;
-use tokio::{runtime::Handle, sync::RwLock};
 
 use swiftide_core::{
     indexing::{IndexingStream, Node},
@@ -38,7 +35,7 @@ impl ScrapingLoader {
 
 impl Loader for ScrapingLoader {
     fn into_stream(mut self) -> IndexingStream {
-        let (tx, rx) = std::sync::mpsc::channel();
+        let (tx, rx) = tokio::sync::mpsc::channel(1000);
         let mut spider_rx = self
             .spider_website
             .subscribe(0)
@@ -58,7 +55,7 @@ impl Loader for ScrapingLoader {
 
                 tracing::debug!(?node, "[Spider] Received node from spider");
 
-                if let Err(error) = tx.send(node) {
+                if let Err(error) = tx.send(node).await {
                     tracing::error!(?error, "[Spider] Failed to send node to stream");
                     break;
                 }
@@ -69,14 +66,14 @@ impl Loader for ScrapingLoader {
 
         let _scrape_thread = tokio::spawn(async move {
             tracing::info!("[Spider] Starting scrape loop");
-            spider_website.scrape().await;
-            spider_website.unsubscribe();
+            // TODO: It would be much nicer if this used `scrape` instead, as it is supposedly
+            // more concurrent
+            spider_website.crawl().await;
             tracing::info!("[Spider] Scrape loop finished");
         });
 
         // NOTE: Handles should stay alive because of rx, but feels a bit fishy
-
-        IndexingStream::iter(rx)
+        rx.into()
     }
 
     fn into_stream_boxed(self: Box<Self>) -> IndexingStream {
@@ -86,11 +83,12 @@ impl Loader for ScrapingLoader {
 
 #[cfg(test)]
 mod tests {
-    use crate::scraping::loader::ScrapingLoader;
-    use futures_util::StreamExt as _;
+    use super::*;
+    use anyhow::Result;
+    use futures_util::StreamExt;
     use swiftide_core::indexing::Loader;
     use wiremock::matchers::{method, path};
-    use wiremock::{Mock, MockServer, ResponseTemplate};
+    use wiremock::{Mock, MockServer, Request, ResponseTemplate};
 
     #[test_log::test(tokio::test(flavor = "multi_thread"))]
     async fn test_scraping_loader_with_wiremock() {
@@ -109,17 +107,58 @@ mod tests {
         let loader = ScrapingLoader::from_url(mock_server.uri());
 
         // Execute the into_stream method
-        let mut stream = loader.into_stream();
+        let stream = loader.into_stream();
 
         // Process the stream to check if we get the expected result
-        while let Some(node) = stream.next().await {
-            tracing::info!(?node, "Received node from stream");
-            // Assert the scraped content against expected content
-            assert_eq!(node.unwrap().chunk, body);
-        }
+        let nodes = stream.collect::<Vec<Result<Node>>>().await;
+
+        assert_eq!(nodes.len(), 1);
+
+        let first_node = nodes.first().unwrap().as_ref().unwrap();
+
+        assert_eq!(first_node.chunk, body);
+    }
+
+    #[test_log::test(tokio::test(flavor = "multi_thread"))]
+    async fn test_scraping_loader_multiple_pages() {
+        // Set up the wiremock server to simulate the remote web server
+        let mock_server = MockServer::start().await;
+
+        // Mocked response for the page we will scrape
+        let body = "<html><body><h1>Test Page</h1><a href=\"/other\">link</a></body></html>";
+        Mock::given(method("GET"))
+            .and(path("/"))
+            .respond_with(ResponseTemplate::new(200).set_body_string(body))
+            .mount(&mock_server)
+            .await;
+
+        let body2 = "<html><body><h1>Test Page 2</h1></body></html>";
+        Mock::given(method("GET"))
+            .and(path("/other"))
+            .respond_with(move |_req: &Request| {
+                std::thread::sleep(std::time::Duration::from_secs(1));
+                ResponseTemplate::new(200).set_body_string(body2)
+            })
+            .mount(&mock_server)
+            .await;
+
+        // Create an instance of ScrapingLoader using the mock server's URL
+        let loader = ScrapingLoader::from_url(mock_server.uri());
+
+        // Execute the into_stream method
+        let stream = loader.into_stream();
+
+        // Process the stream to check if we get the expected result
+        let mut nodes = stream.collect::<Vec<Result<Node>>>().await;
+
+        assert_eq!(nodes.len(), 2);
+
+        let first_node = nodes.pop().unwrap().unwrap();
+
+        assert_eq!(first_node.chunk, body2);
 
-        tracing::info!("Stream finished");
+        let second_node = nodes.pop().unwrap().unwrap();
 
-        drop(stream);
+        assert_eq!(second_node.chunk, body);
     }
 }
