From 3d74de8b5d2e10ab393af36d20069afff3f4a205 Mon Sep 17 00:00:00 2001
From: Alexandr Velikiy <39257464+VargBurz@users.noreply.github.com>
Date: Fri, 13 Dec 2019 12:57:46 +0300
Subject: [PATCH] Unit tests for Bucket class (#811)

---
 analytics/tests/test_bucket.py    | 38 +++++++++++++++++++++++++++++++
 analytics/tests/test_detectors.py | 20 ++++++++++++++++
 2 files changed, 58 insertions(+)
 create mode 100644 analytics/tests/test_bucket.py

diff --git a/analytics/tests/test_bucket.py b/analytics/tests/test_bucket.py
new file mode 100644
index 0000000..8bd138c
--- /dev/null
+++ b/analytics/tests/test_bucket.py
@@ -0,0 +1,38 @@
+import unittest
+import pandas as pd
+import random
+from typing import List
+
+from analytic_types.data_bucket import DataBucket
+from tests.test_dataset import create_list_of_timestamps
+
+class TestBucket(unittest.TestCase):
+
+    def test_receive_data(self):
+        bucket = DataBucket()
+        data_val = list(range(6))
+        timestamp_list = create_list_of_timestamps(len(data_val))
+        for val in data_val:
+            bucket.receive_data(get_pd_dataframe([val], [1523889000000 + val]))
+        for idx, row in bucket.data.iterrows():
+            self.assertEqual(data_val[idx], row['value'])
+            self.assertEqual(timestamp_list[idx], row['timestamp'])
+
+    def test_drop_data(self):
+        bucket = DataBucket()
+        data_val = list(range(10))
+        timestamp_list = create_list_of_timestamps(len(data_val))
+        bucket.receive_data(get_pd_dataframe(data_val, timestamp_list))
+        bucket.drop_data(5)
+        expected_data = data_val[5:]
+        expected_timestamp = timestamp_list[5:]
+        self.assertEqual(expected_data, bucket.data['value'].tolist())
+        self.assertEqual(expected_timestamp, bucket.data['timestamp'].tolist())
+
+if __name__ == '__main__':
+    unittest.main()
+
+def get_pd_dataframe(value: List[int], timestamp: List[int]) -> pd.DataFrame:
+    if len(value) != len(timestamp):
+        raise ValueError(f'len(value) should be equal to len(timestamp)')
+    return pd.DataFrame({ 'value': value, 'timestamp': timestamp })
diff --git a/analytics/tests/test_detectors.py b/analytics/tests/test_detectors.py
index 35951fd..ab4115b 100644
--- a/analytics/tests/test_detectors.py
+++ b/analytics/tests/test_detectors.py
@@ -191,5 +191,25 @@ class TestAnomalyDetector(unittest.TestCase):
         segments_borders = list(map(lambda s: [s.from_timestamp, s.to_timestamp], segments))
         self.assertEqual(segments_borders, [[timestamps[2], timestamps[2]], [timestamps[4], timestamps[8]]])
 
+    def test_consume_data(self):
+        cache =  {
+            'confidence': 2,
+            'alpha': 0.1,
+            'enableBounds': 'ALL',
+            'timeStep': 1
+        }
+        detector = anomaly_detector.AnomalyDetector('test_id')
+
+        detect_result: DetectionResult = None
+        for val in range(22):
+            value = 1 if val != 10 else 5
+            dataframe = pd.DataFrame({'value': [value], 'timestamp': [1523889000000 + val]})
+            dataframe['timestamp'] = pd.to_datetime(dataframe['timestamp'], unit='ms')
+            detect_result = detector.consume_data(dataframe, cache)
+
+        detected_segments = list(map(lambda s: {'from': s.from_timestamp, 'to': s.to_timestamp}, detect_result.segments))
+        result = [{ 'from': 1523889000010, 'to': 1523889000010 }]
+        self.assertEqual(result, detected_segments)
+
 if __name__ == '__main__':
     unittest.main()