From 2dcf1cfb7302901233b46c1d40cce1732033f6c3 Mon Sep 17 00:00:00 2001
From: Avik Basu <ab93@users.noreply.github.com>
Date: Tue, 29 Aug 2023 17:07:28 -0700
Subject: [PATCH] fix: add connector test cases

Signed-off-by: Avik Basu <ab93@users.noreply.github.com>
---
 tests/connectors/__init__.py   |  0
 tests/connectors/test_druid.py | 65 ++++++++++++++++++++++++++++++++++
 tests/connectors/test_redis.py | 17 +++++++++
 3 files changed, 82 insertions(+)
 create mode 100644 tests/connectors/__init__.py
 create mode 100644 tests/connectors/test_druid.py
 create mode 100644 tests/connectors/test_redis.py

diff --git a/tests/connectors/__init__.py b/tests/connectors/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/connectors/test_druid.py b/tests/connectors/test_druid.py
new file mode 100644
index 00000000..f7831950
--- /dev/null
+++ b/tests/connectors/test_druid.py
@@ -0,0 +1,65 @@
+import json
+import unittest
+import datetime
+from unittest.mock import patch, Mock
+import pydruid.query
+from pydruid.client import PyDruid
+from pydruid.utils.aggregators import doublesum
+
+from numalogic.connectors._config import Pivot
+from numalogic.connectors.druid import DruidFetcher
+
+
+def mock_group_by(*_, **__):
+    """Mock group by response from druid."""
+    result = [
+        {
+            "version": "v1",
+            "timestamp": "2023-07-11T01:36:00.000Z",
+            "event": {"count": 5.0, "ciStatus": "success"},
+        },
+        {
+            "version": "v1",
+            "timestamp": "2023-07-11T01:37:00.000Z",
+            "event": {"count": 1.0, "ciStatus": "success"},
+        },
+    ]
+    query = pydruid.query.Query(query_dict={}, query_type="groupBy")
+    query.parse(json.dumps(result))
+    return query
+
+
+class TestDruid(unittest.TestCase):
+    start = None
+    end = None
+    prom = None
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        end = datetime.datetime.now()
+        start = end - datetime.timedelta(hours=36)
+        cls.start = start.timestamp()
+        cls.end = end.timestamp()
+        cls.druid = DruidFetcher(url="http://localhost:8888", endpoint="druid/v2/")
+
+    @patch.object(PyDruid, "groupby", Mock(return_value=mock_group_by()))
+    def test_fetch_data(self):
+        _out = self.druid.fetch_data(
+            filter_keys=["assetId"],
+            filter_values=["5984175597303660107"],
+            dimensions=["ciStatus"],
+            datasource="customer-interaction-metrics",
+            aggregations={"count": doublesum("count")},
+            group_by=["timestamp", "ciStatus"],
+            hours=36,
+            pivot=Pivot(
+                index="timestamp",
+                columns=["ciStatus"],
+                value=["count"],
+            ),
+        )
+        self.assertEqual(_out.shape, (2, 2))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/connectors/test_redis.py b/tests/connectors/test_redis.py
new file mode 100644
index 00000000..a10d18e7
--- /dev/null
+++ b/tests/connectors/test_redis.py
@@ -0,0 +1,17 @@
+import unittest
+from unittest.mock import patch
+
+from redis.sentinel import Sentinel
+import fakeredis
+
+from numalogic.connectors.redis import get_redis_client
+
+server = fakeredis.FakeServer()
+fake_redis_client = fakeredis.FakeStrictRedis(server=server, decode_responses=True)
+
+
+class TestRedisClient(unittest.TestCase):
+    def test_sentinel_redis_client(self):
+        with patch.object(Sentinel, "master_for", return_value=fake_redis_client):
+            r = get_redis_client("hostname", 6379, "pass", "mymaster")
+            self.assertTrue(r.ping())