From 212c03a0e2d1f9e41b92d598c66978d0ceaf6552 Mon Sep 17 00:00:00 2001
From: realaravinth <realaravinth@batsense.net>
Date: Sat, 14 May 2022 18:45:25 +0530
Subject: [PATCH] feat: def interface to get traffic pattern

---
 db/db-core/src/errors.rs | 3 +++
 db/db-core/src/lib.rs    | 9 ++++++++-
 db/db-core/src/tests.rs  | 7 +++++++
 3 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/db/db-core/src/errors.rs b/db/db-core/src/errors.rs
index 7c3bc1a4..702bb0f4 100644
--- a/db/db-core/src/errors.rs
+++ b/db/db-core/src/errors.rs
@@ -44,6 +44,9 @@ pub enum DBError {
     /// Captcha not found
     #[error("Captcha not found")]
     CaptchaNotFound,
+    /// Traffic pattern not found
+    #[error("Traffic pattern not found")]
+    TrafficPatternNotFound,
 }
 
 /// Convenience type alias for grouping driver-specific errors
diff --git a/db/db-core/src/lib.rs b/db/db-core/src/lib.rs
index d80e49c0..1629835c 100644
--- a/db/db-core/src/lib.rs
+++ b/db/db-core/src/lib.rs
@@ -194,9 +194,16 @@ pub trait MCDatabase: std::marker::Send + std::marker::Sync + CloneSPDatabase {
         captcha_key: &str,
         pattern: &TrafficPattern,
     ) -> DBResult<()>;
+
+    /// Get traffic configuration
+    async fn get_traffic_pattern(
+        &self,
+        username: &str,
+        captcha_key: &str,
+    ) -> DBResult<TrafficPattern>;
 }
 
-#[derive(Default, Serialize, Deserialize, Clone, Debug)]
+#[derive(Default, PartialEq, Serialize, Deserialize, Clone, Debug)]
 /// User's traffic pattern; used in generating a captcha configuration
 pub struct TrafficPattern {
     /// average traffic of user's website
diff --git a/db/db-core/src/tests.rs b/db/db-core/src/tests.rs
index 00557afc..b1d185f2 100644
--- a/db/db-core/src/tests.rs
+++ b/db/db-core/src/tests.rs
@@ -140,6 +140,13 @@ pub async fn database_works<'a, T: MCDatabase>(
     // get captcha cooldown duration
     assert_eq!(db.get_captcha_cooldown(c.key).await.unwrap(), c.duration);
 
+    // add traffic pattern
+    db.add_traffic_pattern(p.username, c.key, tp).await.unwrap();
+    assert_eq!(
+        &db.get_traffic_pattern(p.username, c.key).await.unwrap(),
+        tp
+    );
+
     // add captcha levels
     db.add_captcha_levels(p.username, c.key, l).await.unwrap();