add unit stream tests

Signed-off-by: Joel Dice <joel.dice@fermyon.com>

Author: dicej

Cargo.lock

@@ -743,6 +743,20 @@ mod test {
         test_run(&compose(caller, callee).await?).await
     }
 
+    // No-op function; we only test this by composing it in `async_unit_stream_caller`
+    #[allow(
+        dead_code,
+        reason = "here only to make the `assert_test_exists` macro happy"
+    )]
+    fn async_unit_stream_callee() {}
+
+    #[tokio::test]
+    async fn async_unit_stream_caller() -> Result<()> {
+        let caller = &fs::read(test_programs_artifacts::ASYNC_UNIT_STREAM_CALLER_COMPONENT).await?;
+        let callee = &fs::read(test_programs_artifacts::ASYNC_UNIT_STREAM_CALLEE_COMPONENT).await?;
+        test_run(&compose(caller, callee).await?).await
+    }
+
     async fn test_run_bool(component: &[u8], v: bool) -> Result<()> {
         init_logger();
 

crates/misc/component-async-tests/src/lib.rs

@@ -73,6 +73,10 @@ interface run-result {
   run-pass: func() -> result<_, error-context>;
 }
 
+interface unit-stream {
+  run: func(count: u32) -> stream;
+}
+
 world yield-caller {
   import continue;
   import ready;
@@ -153,3 +157,12 @@ world error-context-caller {
   import run-result;
   export run;
 }
+
+world unit-stream-caller {
+  import unit-stream;
+  export run;
+}
+
+world unit-stream-callee {
+  export unit-stream;
+}

crates/misc/component-async-tests/wit/test.wit

@@ -0,0 +1,46 @@
+mod bindings {
+    wit_bindgen::generate!({
+        path: "../misc/component-async-tests/wit",
+        world: "unit-stream-callee",
+        async: {
+            exports: [
+                "local:local/unit-stream#run",
+            ],
+        }
+    });
+
+    use super::Component;
+    export!(Component);
+}
+
+use {
+    bindings::{exports::local::local::unit_stream::Guest, wit_stream},
+    futures::SinkExt,
+    wit_bindgen_rt::async_support::{self, StreamReader},
+};
+
+struct Component;
+
+impl Guest for Component {
+    async fn run(count: u32) -> StreamReader<()> {
+        let (mut tx, rx) = wit_stream::new();
+
+        async_support::spawn(async move {
+            let mut sent = 0;
+            let mut chunk_size = 1;
+            while sent < count {
+                let n = (count - sent).min(chunk_size);
+                tx.send(vec![(); usize::try_from(n).unwrap()])
+                    .await
+                    .unwrap();
+                sent += n;
+                chunk_size *= 2;
+            }
+        });
+
+        rx
+    }
+}
+
+// Unused function; required since this file is built as a `bin`:
+fn main() {}

crates/test-programs/src/bin/async_unit_stream_callee.rs

@@ -0,0 +1,41 @@
+mod bindings {
+    wit_bindgen::generate!({
+        path: "../misc/component-async-tests/wit",
+        world: "unit-stream-caller",
+        async: {
+            imports: [
+                "local:local/unit-stream#run",
+            ],
+            exports: [
+                "local:local/run#run",
+            ],
+        }
+    });
+
+    use super::Component;
+    export!(Component);
+}
+
+use {
+    bindings::{exports::local::local::run::Guest, local::local::unit_stream},
+    futures::StreamExt,
+};
+
+struct Component;
+
+impl Guest for Component {
+    async fn run() {
+        let count = 42;
+        let mut rx = unit_stream::run(count).await;
+
+        let mut received = 0;
+        while let Some(chunk) = rx.next().await {
+            received += chunk.len();
+        }
+
+        assert_eq!(count, u32::try_from(received).unwrap());
+    }
+}
+
+// Unused function; required since this file is built as a `bin`:
+fn main() {}