From 63a03a4fd5cb4ce0ae997036d6237a4fbb4706ea Mon Sep 17 00:00:00 2001
From: Furkan <furkan.bilgin@outlook.com>
Date: Mon, 14 Oct 2024 12:16:43 +0300
Subject: [PATCH] feat: add tests for `DAQJobHandleStats`

---
 src/run_tests.py               |  2 ++
 src/tests/test_handle_stats.py | 53 ++++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+)
 create mode 100644 src/tests/test_handle_stats.py

diff --git a/src/run_tests.py b/src/run_tests.py
index 13274a0..7995851 100644
--- a/src/run_tests.py
+++ b/src/run_tests.py
@@ -1,6 +1,7 @@
 import unittest
 
 from tests.test_csv import TestDAQJobStoreCSV
+from tests.test_handle_stats import TestDAQJobHandleStats
 from tests.test_main import TestMain
 from tests.test_n1081b import TestDAQJobN1081B
 from tests.test_slack import TestDAQJobAlertSlack
@@ -12,6 +13,7 @@ def run_tests():
     test_suite.addTests(loader.loadTestsFromTestCase(TestDAQJobStoreCSV))
     test_suite.addTests(loader.loadTestsFromTestCase(TestDAQJobN1081B))
     test_suite.addTests(loader.loadTestsFromTestCase(TestMain))
+    test_suite.addTests(loader.loadTestsFromTestCase(TestDAQJobHandleStats))
     test_suite.addTests(loader.loadTestsFromTestCase(TestDAQJobAlertSlack))
     return test_suite
 
diff --git a/src/tests/test_handle_stats.py b/src/tests/test_handle_stats.py
new file mode 100644
index 0000000..9ccbb04
--- /dev/null
+++ b/src/tests/test_handle_stats.py
@@ -0,0 +1,53 @@
+import unittest
+from datetime import datetime
+from typing import cast
+from unittest.mock import MagicMock, patch
+
+from daq.base import DAQJob
+from daq.jobs.handle_stats import (
+    DAQJobHandleStats,
+    DAQJobMessageStats,
+    DAQJobStatsRecord,
+)
+from daq.models import DAQJobStats
+
+
+class TestDAQJobHandleStats(unittest.TestCase):
+    def setUp(self):
+        self.config = MagicMock()
+        self.daq_job_handle_stats = DAQJobHandleStats(config=self.config)
+        self.daq_job_handle_stats.message_out = MagicMock()
+
+    def test_handle_message_success(self):
+        message = DAQJobMessageStats(
+            stats={
+                type(cast(DAQJob, MagicMock())): DAQJobStats(
+                    message_in_stats=DAQJobStatsRecord(
+                        last_updated=datetime.now(), count=5
+                    ),
+                    message_out_stats=DAQJobStatsRecord(
+                        last_updated=datetime.now(), count=3
+                    ),
+                    restart_stats=DAQJobStatsRecord(
+                        last_updated=datetime.now(), count=1
+                    ),
+                )
+            }
+        )
+
+        result = self.daq_job_handle_stats.handle_message(message)
+
+        self.assertTrue(result)
+        self.daq_job_handle_stats.message_out.put.assert_called_once()
+
+    def test_handle_message_failure(self):
+        message = MagicMock()
+        with patch("daq.jobs.handle_stats.DAQJob.handle_message", return_value=False):
+            result = self.daq_job_handle_stats.handle_message(message)
+
+        self.assertFalse(result)
+        self.daq_job_handle_stats.message_out.put.assert_not_called()
+
+
+if __name__ == "__main__":
+    unittest.main()