fix test loggers, enable peers function

Author: erhant

compute/src/node.rs

@@ -21,19 +21,6 @@ use crate::{
 /// Number of seconds between refreshing the Kademlia DHT.
 const PEER_REFRESH_INTERVAL_SECS: u64 = 30;
 
-/// **Dria Compute Node**
-///
-/// Internally, the node will create a new P2P client with the given secret key.
-/// This P2P client, although created synchronously, requires a tokio runtime.
-///
-/// ### Example
-///
-/// ```rs
-/// let config = DriaComputeNodeConfig::new();
-/// let mut node = DriaComputeNode::new(config, CancellationToken::new())?;
-/// node.check_services().await?;
-/// node.launch().await?;
-/// ```
 pub struct DriaComputeNode {
     pub config: DriaComputeNodeConfig,
     pub p2p: DriaP2PCommander,
@@ -44,10 +31,13 @@ pub struct DriaComputeNode {
 }
 
 impl DriaComputeNode {
+    /// Creates a new `DriaComputeNode` with the given configuration and cancellation token.
+    ///
+    /// Returns the node instance and p2p client together. P2p MUST be run in a separate task before this node is used at all.
     pub async fn new(
         config: DriaComputeNodeConfig,
         cancellation: CancellationToken,
-    ) -> Result<(Self, DriaP2PClient)> {
+    ) -> Result<(DriaComputeNode, DriaP2PClient)> {
         // create the keypair from secret key
         let keypair = secret_to_keypair(&config.secret_key);
 
@@ -139,16 +129,11 @@ impl DriaComputeNode {
         Ok(())
     }
 
-    /// Returns the list of connected peers.
-    // #[inline(always)]
-    // pub fn peers(
-    //     &self,
-    // ) -> Vec<(
-    //     &dkn_p2p::libp2p_identity::PeerId,
-    //     Vec<&gossipsub::TopicHash>,
-    // )> {
-    //     self.p2p.peers()
-    // }
+    /// Returns the list of connected peers, `mesh` and `all`.
+    #[inline]
+    pub async fn peers(&self) -> Result<(Vec<PeerId>, Vec<PeerId>)> {
+        self.p2p.peers().await
+    }
 
     /// Launches the main loop of the compute node.
     /// This method is not expected to return until cancellation occurs.
@@ -285,16 +270,24 @@ impl DriaComputeNode {
         self.unsubscribe(WorkflowHandler::LISTEN_TOPIC).await?;
         self.unsubscribe(WorkflowHandler::RESPONSE_TOPIC).await?;
 
+        // shutdown channels
+        self.shutdown().await?;
+
+        Ok(())
+    }
+
+    /// Shutdown channels between p2p and yourself.
+    pub async fn shutdown(&mut self) -> Result<()> {
         // send shutdown signal
+        log::debug!("Sending shutdown command to p2p client.");
         self.p2p.shutdown().await?;
 
-        // close msg channel
+        // close message channel
         log::debug!("Closing message channel.");
         self.msg_rx.close();
 
         Ok(())
     }
-
     /// Parses a given raw Gossipsub message to a prepared P2PMessage object.
     /// This prepared message includes the topic, payload, version and timestamp.
     ///
@@ -322,34 +315,47 @@ impl DriaComputeNode {
 
 #[cfg(test)]
 mod tests {
-    // use super::*;
-    // use std::env;
-
-    // #[tokio::test]
-    // #[ignore = "run this manually"]
-    // async fn test_publish_message() {
-    //     env::set_var("RUST_LOG", "info");
-    //     let _ = env_logger::try_init();
-
-    //     // create node
-    //     let cancellation = CancellationToken::new();
-    //     let mut node = DriaComputeNode::new(DriaComputeNodeConfig::default(), cancellation.clone())
-    //         .await
-    //         .expect("should create node");
-
-    //     // launch & wait for a while for connections
-    //     log::info!("Waiting a bit for peer setup.");
-    //     tokio::select! {
-    //         _ = node.launch() => (),
-    //         _ = tokio::time::sleep(tokio::time::Duration::from_secs(20)) => cancellation.cancel(),
-    //     }
-    //     log::info!("Connected Peers:\n{:#?}", node.peers());
-
-    //     // publish a dummy message
-    //     let topic = "foo";
-    //     let message = DKNMessage::new("hello from the other side", topic);
-    //     node.subscribe(topic).expect("should subscribe");
-    //     node.publish(message).expect("should publish");
-    //     node.unsubscribe(topic).expect("should unsubscribe");
-    // }
+    use super::*;
+    use std::env;
+
+    #[tokio::test]
+    #[ignore = "run this manually"]
+    async fn test_publish_message() -> eyre::Result<()> {
+        env::set_var("RUST_LOG", "none,dkn_compute=debug,dkn_p2p=debug");
+        let _ = env_logger::builder().is_test(true).try_init();
+
+        // create node
+        let cancellation = CancellationToken::new();
+        let (mut node, p2p) =
+            DriaComputeNode::new(DriaComputeNodeConfig::default(), cancellation.clone())
+                .await
+                .expect("should create node");
+
+        // spawn p2p task
+        let p2p_task = tokio::spawn(async move { p2p.run().await });
+
+        // launch & wait for a while for connections
+        log::info!("Waiting a bit for peer setup.");
+        tokio::select! {
+            _ = node.launch() => (),
+            _ = tokio::time::sleep(tokio::time::Duration::from_secs(20)) => cancellation.cancel(),
+        }
+        log::info!("Connected Peers:\n{:#?}", node.peers().await?);
+
+        // publish a dummy message
+        let topic = "foo";
+        let message = DKNMessage::new("hello from the other side", topic);
+        node.subscribe(topic).await.expect("should subscribe");
+        node.publish(message).await.expect("should publish");
+        node.unsubscribe(topic).await.expect("should unsubscribe");
+
+        // close everything
+        log::info!("Shutting down node.");
+        node.p2p.shutdown().await?;
+
+        // wait for task handle
+        p2p_task.await?;
+
+        Ok(())
+    }
 }

p2p/src/commands.rs

@@ -233,7 +233,7 @@ impl DriaP2PCommander {
     }
 
     /// Get peers counts (mesh & all) of the GossipSub pool.
-    /// Returns a tuple of the mesh peer count and all peer count.
+    /// Returns a tuple of the mesh peers count and all peers count.
     pub async fn peer_counts(&self) -> Result<(usize, usize)> {
         let (sender, receiver) = oneshot::channel();
 

p2p/tests/listen_test.rs

@@ -4,47 +4,40 @@ use libp2p::Multiaddr;
 use libp2p_identity::Keypair;
 use std::{env, str::FromStr};
 
-const TOPIC: &str = "pong";
-const LOG_LEVEL: &str = "none,listen_test=debug,dkn_p2p=debug";
-
 // FIXME: not working!!!
 #[tokio::test]
-#[ignore = "run manually with logs"]
+#[ignore = "run this manually"]
 async fn test_listen_topic_once() -> Result<()> {
-    env::set_var("RUST_LOG", LOG_LEVEL);
-    let _ = env_logger::builder().is_test(true).try_init();
+    const TOPIC: &str = "pong";
 
-    // setup client
-    let keypair = Keypair::generate_secp256k1();
-    let addr = Multiaddr::from_str("/ip4/0.0.0.0/tcp/4001")?;
-    let bootstraps = vec![Multiaddr::from_str(
-        "/ip4/44.206.245.139/tcp/4001/p2p/16Uiu2HAm4q3LZU2T9kgjKK4ysy6KZYKLq8KiXQyae4RHdF7uqSt4",
-    )?];
-    let relays = vec![Multiaddr::from_str(
-        "/ip4/34.201.33.141/tcp/4001/p2p/16Uiu2HAkuXiV2CQkC9eJgU6cMnJ9SMARa85FZ6miTkvn5fuHNufa",
-    )?];
-    let protocol = DriaP2PProtocol::new_major_minor("dria");
+    env::set_var("RUST_LOG", "none,listen_test=debug,dkn_p2p=debug");
+    let _ = env_logger::builder().is_test(true).try_init();
 
     // spawn P2P client in another task
     let (client, mut commander, mut msg_rx) = DriaP2PClient::new(
-        keypair,
-        addr,
-        bootstraps.into_iter(),
-        relays.into_iter(),
+        Keypair::generate_secp256k1(),
+        Multiaddr::from_str("/ip4/0.0.0.0/tcp/4001")?,
+        vec![Multiaddr::from_str(
+            "/ip4/44.206.245.139/tcp/4001/p2p/16Uiu2HAm4q3LZU2T9kgjKK4ysy6KZYKLq8KiXQyae4RHdF7uqSt4",
+        )?].into_iter(),
+        vec![Multiaddr::from_str(
+            "/ip4/34.201.33.141/tcp/4001/p2p/16Uiu2HAkuXiV2CQkC9eJgU6cMnJ9SMARa85FZ6miTkvn5fuHNufa",
+        )?]
+        .into_iter(),
         vec![].into_iter(),
-        protocol,
+        DriaP2PProtocol::new_major_minor("dria"),
     )
     .expect("could not create p2p client");
 
+    // spawn task
+    let p2p_task = tokio::spawn(async move { client.run().await });
+
     // subscribe to the given topic
     commander
         .subscribe(TOPIC)
         .await
         .expect("could not subscribe");
 
-    // spawn task
-    let p2p_task = tokio::spawn(async move { client.run().await });
-
     // wait for a single gossipsub message on this topic
     log::info!("Waiting for messages...");
     let message = msg_rx.recv().await;

workflows/src/providers/gemini.rs

@@ -185,7 +185,7 @@ mod tests {
         let _ = dotenvy::dotenv(); // read api key
         assert!(env::var(ENV_VAR_NAME).is_ok(), "should have api key");
         env::set_var("RUST_LOG", "none,dkn_workflows=debug");
-        let _ = env_logger::try_init();
+        let _ = env_logger::builder().is_test(true).try_init();
 
         let models = vec![
             Model::Gemini10Pro,

workflows/src/providers/openai.rs

@@ -172,7 +172,7 @@ mod tests {
         let _ = dotenvy::dotenv(); // read api key
         assert!(env::var(ENV_VAR_NAME).is_ok(), "should have api key");
         env::set_var("RUST_LOG", "none,dkn_workflows=debug");
-        let _ = env_logger::try_init();
+        let _ = env_logger::builder().is_test(true).try_init();
 
         let models = vec![Model::GPT4Turbo, Model::GPT4o, Model::GPT4oMini];
         let config = OpenAIConfig::new();

workflows/src/providers/openrouter.rs

@@ -120,7 +120,7 @@ mod tests {
         let _ = dotenvy::dotenv(); // read api key
         assert!(env::var(ENV_VAR_NAME).is_ok(), "should have api key");
         env::set_var("RUST_LOG", "none,dkn_workflows=debug");
-        let _ = env_logger::try_init();
+        let _ = env_logger::builder().is_test(true).try_init();
 
         let models = vec![Model::GPT4Turbo, Model::GPT4o, Model::GPT4oMini];
         let config = OpenRouterConfig::new();

workflows/tests/models_test.rs

@@ -8,7 +8,7 @@ const LOG_LEVEL: &str = "none,dkn_workflows=debug";
 #[ignore = "requires Ollama"]
 async fn test_ollama_check() -> Result<()> {
     env::set_var("RUST_LOG", LOG_LEVEL);
-    let _ = env_logger::try_init();
+    let _ = env_logger::builder().is_test(true).try_init();
 
     let models = vec![Model::Phi3_5Mini];
     let mut model_config = DriaWorkflowsConfig::new(models);
@@ -27,7 +27,7 @@ async fn test_ollama_check() -> Result<()> {
 async fn test_openai_check() -> Result<()> {
     let _ = dotenvy::dotenv(); // read api key
     env::set_var("RUST_LOG", LOG_LEVEL);
-    let _ = env_logger::try_init();
+    let _ = env_logger::builder().is_test(true).try_init();
 
     let models = vec![Model::GPT4Turbo];
     let mut model_config = DriaWorkflowsConfig::new(models);