feat: add out_dir to DAQJobStoreCSV and update tests

ENRG-tr · Nov 6, 2024 · 759234d · 759234d
1 parent d4bad71
commit 759234d
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 8 deletions.
diff --git a/src/daq/jobs/store/csv.py b/src/daq/jobs/store/csv.py
@@ -25,7 +25,7 @@ class DAQJobStoreConfigCSV(DAQJobStoreConfig):
 
 @dataclass
 class DAQJobStoreCSVConfig(DAQJobConfig):
-    pass
+    out_dir: str = "out/"
 
 
 @dataclass
@@ -53,6 +53,8 @@ def handle_message(self, message: DAQJobMessageStore) -> bool:
         file_path = modify_file_path(
             store_config.file_path, store_config.add_date, message.prefix
         )
+        file_path = os.path.join(self.config.out_dir, file_path)
+
         file, new_file = self._open_csv_file(file_path, store_config.overwrite)
         if file.overwrite:
             file.write_queue.clear()

diff --git a/src/tests/test_csv.py b/src/tests/test_csv.py
@@ -14,7 +14,7 @@
 
 class TestDAQJobStoreCSV(unittest.TestCase):
     def setUp(self):
-        self.config = MagicMock()
+        self.config = MagicMock(out_dir="out/")
         self.store = DAQJobStoreCSV(self.config)
 
     @patch("daq.jobs.store.csv.modify_file_path", return_value="test.csv")
@@ -35,9 +35,9 @@ def test_handle_message_new_file(
         self.store.handle_message(message)
 
         mock_add_date.assert_called_once_with("test.csv", True, None)
-        mock_open.assert_called_once_with("test.csv", "a", newline="")
-        self.assertIn("test.csv", self.store._open_csv_files)
-        file = self.store._open_csv_files["test.csv"]
+        mock_open.assert_called_once_with("out/test.csv", "a", newline="")
+        self.assertIn("out/test.csv", self.store._open_csv_files)
+        file = self.store._open_csv_files["out/test.csv"]
         self.assertEqual(len(file.write_queue), 3)  # 1 header + 2 rows
 
     @patch("daq.jobs.store.csv.modify_file_path", return_value="test.csv")
@@ -55,9 +55,9 @@ def test_handle_message_existing_file(self, mock_exists, mock_open, mock_add_dat
         self.store.handle_message(message)
 
         mock_add_date.assert_called_once_with("test.csv", True, None)
-        mock_open.assert_called_once_with("test.csv", "a", newline="")
-        self.assertIn("test.csv", self.store._open_csv_files)
-        file = self.store._open_csv_files["test.csv"]
+        mock_open.assert_called_once_with("out/test.csv", "a", newline="")
+        self.assertIn("out/test.csv", self.store._open_csv_files)
+        file = self.store._open_csv_files["out/test.csv"]
         self.assertEqual(len(file.write_queue), 2)  # 2 rows only, no header
 
     def test_flush(self):