From 2eae205d550d327b6fee2286d91eae35c83e56fb Mon Sep 17 00:00:00 2001
From: Colin Roberts <colin@primitive.xyz>
Date: Fri, 16 Feb 2024 11:48:33 -0700
Subject: [PATCH] test: no-stream

---
 engine/tests/common.rs              | 10 ++++++++--
 engine/tests/machine_integration.rs | 28 ++++++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 2 deletions(-)
 create mode 100644 engine/tests/machine_integration.rs

diff --git a/engine/tests/common.rs b/engine/tests/common.rs
index 883ef2ef..9416f270 100644
--- a/engine/tests/common.rs
+++ b/engine/tests/common.rs
@@ -8,6 +8,12 @@ use arbiter_engine::{
 };
 use serde::{Deserialize, Serialize};
 
+#[allow(unused)]
+fn trace() {
+    std::env::set_var("RUST_LOG", "trace");
+    tracing_subscriber::fmt::init();
+}
+
 fn default_max_count() -> Option<u64> {
     Some(3)
 }
@@ -53,12 +59,12 @@ impl Behavior<Message> for TimedMessage {
         &mut self,
         _client: Arc<ArbiterMiddleware>,
         messager: Messager,
-    ) -> Result<EventStream<Message>> {
+    ) -> Result<Option<EventStream<Message>>> {
         if let Some(startup_message) = &self.startup_message {
             messager.send(To::All, startup_message).await?;
         }
         self.messager = Some(messager.clone());
-        Ok(messager.stream()?)
+        Ok(Some(messager.stream()?))
     }
 
     async fn process(&mut self, event: Message) -> Result<ControlFlow> {
diff --git a/engine/tests/machine_integration.rs b/engine/tests/machine_integration.rs
new file mode 100644
index 00000000..3538daea
--- /dev/null
+++ b/engine/tests/machine_integration.rs
@@ -0,0 +1,28 @@
+use arbiter_engine::{agent::Agent, world::World};
+
+include!("common.rs");
+
+#[derive(Debug, Deserialize, Serialize)]
+struct MockBehavior;
+
+#[async_trait::async_trait]
+impl Behavior<()> for MockBehavior {
+    async fn startup(
+        &mut self,
+        _client: Arc<ArbiterMiddleware>,
+        _messager: Messager,
+    ) -> Result<Option<EventStream<()>>> {
+        Ok(None)
+    }
+}
+
+#[tokio::test]
+async fn behavior_no_stream() {
+    trace();
+    let mut world = World::new("test");
+    let behavior = MockBehavior;
+    let agent = Agent::builder("agent").with_behavior(behavior);
+    world.add_agent(agent);
+
+    world.run().await.unwrap();
+}