From a4929922bf1cac3f67014bdc36a857cfdba017c1 Mon Sep 17 00:00:00 2001
From: realaravinth <realaravinth@batsense.net>
Date: Mon, 15 Mar 2021 19:37:42 +0530
Subject: [PATCH] GetCurrentVisitor count

---
 examples/simple.rs |  2 --
 src/mcaptcha.rs    | 26 ++++++++++++++++++++++++++
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/examples/simple.rs b/examples/simple.rs
index bb196b0..a386a73 100644
--- a/examples/simple.rs
+++ b/examples/simple.rs
@@ -67,8 +67,6 @@ async fn main() -> std::io::Result<()> {
         .build()
         .unwrap();
 
-    p
-
     // create and start MCaptcha actor that uses the above defense configuration
     // This is what manages the difficulty factor of sites that an mCaptcha protects
     let mcaptcha = MCaptchaBuilder::default()
diff --git a/src/mcaptcha.rs b/src/mcaptcha.rs
index 625f9fb..79e930c 100644
--- a/src/mcaptcha.rs
+++ b/src/mcaptcha.rs
@@ -229,6 +229,19 @@ impl Handler<AddVisitor> for MCaptcha {
     }
 }
 
+/// Message to get the visitor count
+#[derive(Message)]
+#[rtype(result = "u32")]
+pub struct GetCurrentVisitorCount;
+
+impl Handler<GetCurrentVisitorCount> for MCaptcha {
+    type Result = MessageResult<GetCurrentVisitorCount>;
+
+    fn handle(&mut self, _: GetCurrentVisitorCount, _ctx: &mut Self::Context) -> Self::Result {
+        MessageResult(self.visitor_threshold)
+    }
+}
+
 #[cfg(test)]
 pub mod tests {
     use super::*;
@@ -331,4 +344,17 @@ pub mod tests {
             Some(CaptchaError::PleaseSetValue("duration".into()))
         );
     }
+
+    #[actix_rt::test]
+    async fn get_current_visitor_count_works() {
+        let addr: MyActor = get_counter().start();
+
+        addr.send(AddVisitor).await.unwrap();
+        addr.send(AddVisitor).await.unwrap();
+        addr.send(AddVisitor).await.unwrap();
+        addr.send(AddVisitor).await.unwrap();
+        let count = addr.send(GetCurrentVisitorCount).await.unwrap();
+
+        assert_eq!(count, 4);
+    }
 }