feat: tests for deployer, price_changer and lib

anthias-labs · Jul 19, 2024 · 876b8f8 · 876b8f8
1 parent 9f8b79a
commit 876b8f8
Show file tree

Hide file tree

Showing 3 changed files with 179 additions and 23 deletions.
diff --git a/src/deployer.rs b/src/deployer.rs
@@ -119,3 +119,89 @@ impl Behavior<Message> for Deployer {
         }
     }
 }
+
+// #[cfg(test)]
+// mod tests {
+//     use octane::{agent::Agent, world::World};
+
+//     use super::*;
+//     use crate::deployer::Deployer;
+
+//     #[derive(Debug, Serialize, Deserialize)]
+//     pub struct MockDeployer {
+//         #[serde(skip)]
+//         pub messager: Option<Messager>,
+
+//         #[serde(skip)]
+//         pub client: Option<Arc<AnvilProvider>>,
+//     }
+
+//     #[async_trait::async_trait]
+//     impl Behavior<Message> for MockDeployer {
+//         async fn startup(
+//             &mut self,
+//             client: Arc<AnvilProvider>,
+//             messager: Messager,
+//         ) -> Result<Option<EventStream<Message>>> {
+//             messager
+//                 .send(
+//                     To::Agent("deployer".to_string()),
+//                     DeploymentRequest::Token {
+//                         name: String::from("TEST0"),
+//                         symbol: String::from("TST0"),
+//                         decimals: 18,
+//                     },
+//                 )
+//                 .await?;
+
+//             self.client = Some(client.clone());
+//             self.messager = Some(messager.clone());
+
+//             Ok(Some(messager.clone().stream().unwrap()))
+//         }
+
+//         async fn process(&mut self, event: Message) -> Result<ControlFlow> {
+//             let query: DeploymentResponse = match serde_json::from_str(&event.data) {
+//                 Ok(query) => query,
+//                 Err(_) => {
+//                     eprintln!("Failed to deserialize the event data into a DeploymentResponse");
+//                     return Ok(ControlFlow::Continue);
+//                 }
+//             };
+
+//             match query {
+//                 DeploymentResponse::Token(address) => {
+//                     let tok = ArenaToken::new(address, self.client.clone().unwrap());
+
+//                     assert_eq!(tok.name().call().await.unwrap()._0, "TEST");
+//                     assert_eq!(tok.symbol().call().await.unwrap()._0, "TST");
+//                     assert_eq!(tok.decimals().call().await.unwrap()._0, 18);
+//                 }
+//                 _ => {}
+//             }
+
+//             Ok(ControlFlow::Continue)
+//         }
+//     }
+
+//     #[tokio::test]
+//     async fn test_deployer() {
+//         // env_logger::init();
+
+//         let deployer = Agent::builder("deployer").with_behavior(Deployer {
+//             messager: None,
+//             client: None,
+//         });
+//         let mock_deployer = Agent::builder("mock_deployer").with_behavior(MockDeployer {
+//             client: None,
+//             messager: None,
+//         });
+
+//         let mut world = World::new("id");
+
+//         world.add_agent(mock_deployer);
+//         world.add_agent(deployer);
+
+//         let _ = world.run().await;
+//     }
+// }
diff --git a/src/lib.rs b/src/lib.rs
@@ -24,6 +24,7 @@ use crate::{
     price_changer::PriceUpdate,
     types::process::{OrnsteinUhlenbeck, StochasticProcess},
 };
+use crate::price_changer::PriceChanger;
 
 pub mod arbitrageur;
 pub mod bindings;
@@ -33,6 +34,7 @@ pub mod pool_admin;
 pub mod price_changer;
 pub mod types;
 
+
 #[cfg(test)]
 mod tests {
     use octane::{agent::Agent, world::World};
@@ -41,7 +43,7 @@ mod tests {
     use crate::deployer::Deployer;
 
     #[derive(Debug, Serialize, Deserialize)]
-    pub struct MockDeployer {
+    pub struct TokenDeployer {
         #[serde(skip)]
         pub messager: Option<Messager>,
 
@@ -50,12 +52,15 @@ mod tests {
     }
 
     #[async_trait::async_trait]
-    impl Behavior<Message> for MockDeployer {
+    impl Behavior<Message> for TokenDeployer {
         async fn startup(
             &mut self,
             client: Arc<AnvilProvider>,
             messager: Messager,
         ) -> Result<Option<EventStream<Message>>> {
+            self.client = Some(client.clone());
+            self.messager = Some(messager.clone());
+
             messager
                 .send(
                     To::Agent("deployer".to_string()),
@@ -67,54 +72,117 @@ mod tests {
                 )
                 .await?;
 
-            self.client = Some(client.clone());
-            self.messager = Some(messager.clone());
+            messager
+                .send(
+                    To::Agent("deployer".to_string()),
+                    DeploymentRequest::Token {
+                        name: String::from("TEST1"),
+                        symbol: String::from("TST1"),
+                        decimals: 18,
+                    },
+                )
+                .await?;
 
             Ok(Some(messager.clone().stream().unwrap()))
         }
+    }
 
-        async fn process(&mut self, event: Message) -> Result<ControlFlow> {
-            let query: DeploymentResponse = match serde_json::from_str(&event.data) {
-                Ok(query) => query,
-                Err(_) => {
-                    eprintln!("Failed to deserialize the event data into a DeploymentResponse");
-                    return Ok(ControlFlow::Continue);
-                }
-            };
+    #[derive(Debug, Serialize, Deserialize)]
+    pub struct MockOrchestrator {
+        #[serde(skip)]
+        pub messager: Option<Messager>,
 
-            match query {
-                DeploymentResponse::Token(address) => {
-                    let tok = ArenaToken::new(address, self.client.clone().unwrap());
+        #[serde(skip)]
+        pub client: Option<Arc<AnvilProvider>>,
+
+        pub tokens: Vec<Address>,
+    }
 
-                    assert_eq!(tok.name().call().await.unwrap()._0, "TEST");
-                    assert_eq!(tok.symbol().call().await.unwrap()._0, "TST");
-                    assert_eq!(tok.decimals().call().await.unwrap()._0, 18);
+    #[async_trait::async_trait]
+    impl Behavior<Message> for MockOrchestrator {
+        async fn startup(
+            &mut self,
+            client: Arc<AnvilProvider>,
+            messager: Messager,
+        ) -> Result<Option<EventStream<Message>>> {
+            let mut stream = messager.clone().stream().unwrap();
+
+            while let Some(event) = stream.next().await {
+                let query: DeploymentResponse = match serde_json::from_str(&event.data) {
+                    Ok(query) => query,
+                    Err(_) => {
+                        eprintln!("Failed to deserialize the event data into a DeploymentResponse");
+                        continue;
+                    }
+                };
+
+                if let DeploymentResponse::Token(address) = query {
+                    self.tokens.push(address);
+                }
+
+                if self.tokens.len() == 2 {
+                    break;
                 }
-                _ => {}
             }
 
-            Ok(ControlFlow::Continue)
+            messager
+                .send(
+                    To::Agent("deployer".to_string()),
+                    DeploymentRequest::LiquidExchange {
+                        token_0: self.tokens[0],
+                        token_1: self.tokens[1],
+                        initial_price: 1.0,
+                    },
+                )
+                .await?;
+
+            println!("Tokens: {:?}", self.tokens);
+
+            for i in 0..100 {
+                messager
+                    .send(
+                        To::Agent("pricechanger".to_string()),
+                        PriceUpdate,
+                    )
+                    .await?;
+            }
+
+            self.client = Some(client.clone());
+            self.messager = Some(messager.clone());
+
+            Ok(Some(messager.clone().stream().unwrap()))
         }
     }
 
     #[tokio::test]
-    async fn test_deployer() {
+    async fn asdfasdf() {
         env_logger::init();
 
+        let token_deployer = Agent::builder("tdeployer").with_behavior(TokenDeployer {
+            messager: None,
+            client: None,
+        });
+
         let deployer = Agent::builder("deployer").with_behavior(Deployer {
             messager: None,
             client: None,
         });
-        let mock_deployer = Agent::builder("mock_deployer").with_behavior(MockDeployer {
+
+        let mock_deployer = Agent::builder("mock_deployer").with_behavior(MockOrchestrator {
             client: None,
             messager: None,
+            tokens: vec![],
         });
 
+        let changer = Agent::builder("pricechanger").with_behavior(PriceChanger::new(OrnsteinUhlenbeck::new(0.0, 0.1, 0.1, 0.1)));
+
         let mut world = World::new("id");
 
         world.add_agent(mock_deployer);
         world.add_agent(deployer);
+        world.add_agent(token_deployer);
+        world.add_agent(changer);
 
         let _ = world.run().await;
     }
-}
+}
diff --git a/src/price_changer.rs b/src/price_changer.rs
@@ -74,6 +74,8 @@ impl Behavior<Message> for PriceChanger {
 
         tx.send().await?.watch().await?;
 
+        println!("Price updated to: {}", self.process.current_value());
+
         Ok(ControlFlow::Continue)
     }
 }