taskbuffer/tests/key_reads_tests.py

0001 # Description: Unit tests for the reads of jobs, tasks and datasets by key. No database is used
0002 import datetime
0003 import unittest
0004 from typing import Any
0005 from unittest import mock
0006
0007 from pandaserver.taskbuffer import task_split_rules
0008 from pandaserver.taskbuffer.DatasetSpec import DatasetSpec
0009 from pandaserver.taskbuffer.db_proxy_mods import job_standalone_module
0010 from pandaserver.taskbuffer.db_proxy_mods.job_standalone_module import (
0011     JobStandaloneModule,
0012 )
0013 from pandaserver.taskbuffer.db_proxy_mods.misc_standalone_module import (
0014     MiscStandaloneModule,
0015 )
0016 from pandaserver.taskbuffer.db_proxy_mods.task_utils_module import TaskUtilsModule
0017 from pandaserver.taskbuffer.JediTaskSpec import JediTaskSpec
0018 from pandaserver.taskbuffer.JobSpec import JobSpec
0019 from pandaserver.taskbuffer.TaskBuffer import TaskBuffer
0020
0021 LOST = "ORA-03113: end-of-file on communication channel"
0022
0023
0024 class TestPeekRaiseError(unittest.TestCase):
0025     def setUp(self) -> None:
0026         self.proxy = mock.MagicMock()
0027         self.proxy.cur.execute.side_effect = Exception(LOST)
0028         patcher = mock.patch.object(job_standalone_module.time, "sleep")
0029         self.sleep = patcher.start()
0030         self.addCleanup(patcher.stop)
0031
0032     def test_peek_job_raises_at_once(self) -> None:
0033         with self.assertRaisesRegex(Exception, "ORA-03113"):
0034             JobStandaloneModule.peekJob(self.proxy, 7000000000, True, True, True, True, forAnal=True, raise_error=True)
0035         self.assertEqual(self.proxy.cur.execute.call_count, 1)
0036         self.sleep.assert_not_called()
0037
0038     def test_peek_job_as_before_without(self) -> None:
0039         self.assertIsNone(JobStandaloneModule.peekJob(self.proxy, 7000000000, True, True, True, True, forAnal=True))
0040         self.assertEqual(self.proxy.cur.execute.call_count, 3)
0041         job = JobStandaloneModule.peekJob(self.proxy, 7000000000, True, True, True, True)
0042         self.assertIsInstance(job, JobSpec)
0043         self.assertEqual(job.jobStatus, "unknown")  # type: ignore[union-attr]
0044
0045     def test_peek_job_log(self) -> None:
0046         with self.assertRaisesRegex(Exception, "ORA-03113"):
0047             JobStandaloneModule.peekJobLog(self.proxy, 7000000000, 90, raise_error=True)
0048         self.assertIsNone(JobStandaloneModule.peekJobLog(self.proxy, 7000000000, 90))
0049
0050
0051 class TestGetFullJobStatus(unittest.TestCase):
0052     """Job 1 is live, 2 is nowhere, 3 a live table cannot be read for, and 4 the archive cannot be read for."""
0053
0054     def setUp(self) -> None:
0055         self.job = JobSpec()
0056         self.proxy = mock.MagicMock()
0057
0058         def peek_job(job_id: int, *args: Any, raise_error: bool = False) -> JobSpec | None:
0059             if job_id == 3 and raise_error:
0060                 raise Exception(LOST)
0061             return self.job if job_id == 1 else None
0062
0063         def peek_job_log(job_id: int, days: int | None, raise_error: bool = False) -> JobSpec | None:
0064             if job_id == 4 and raise_error:
0065                 raise Exception("ORA-01013: user requested cancel of current operation")
0066             return None
0067
0068         self.proxy.peekJob.side_effect = peek_job
0069         self.proxy.peekJobLog.side_effect = peek_job_log
0070         self.task_buffer = TaskBuffer.__new__(TaskBuffer)
0071         self.task_buffer.proxyPool = mock.MagicMock()
0072         self.task_buffer.proxyPool.get.return_value.__enter__.return_value = self.proxy
0073
0074     def test_errors_reported(self) -> None:
0075         jobs = self.task_buffer.getFullJobStatus([1, 2, 3, 4], report_errors=True)
0076         self.assertEqual(jobs, [self.job, None, {"error": "ORA-03113"}, {"error": "ORA-01013"}])
0077         # the job the live tables could not be read for is not looked for in the archive
0078         self.assertEqual([call.args[0] for call in self.proxy.peekJobLog.call_args_list], [2, 4])
0079
0080     def test_errors_as_not_found_by_default(self) -> None:
0081         self.assertEqual(self.task_buffer.getFullJobStatus([1, 2, 3, 4]), [self.job, None, None, None])
0082         self.assertEqual({call.kwargs["raise_error"] for call in self.proxy.peekJob.call_args_list}, {False})
0083
0084     def test_unexpected_error_raised_by_default(self) -> None:
0085         self.proxy.peekJob.side_effect = ValueError("broken")
0086         with self.assertRaises(ValueError):
0087             self.task_buffer.getFullJobStatus([1])
0088
0089
0090 class TestGetDatasetsByName(unittest.TestCase):
0091     AS_OF = datetime.datetime(2026, 10, 14, 9, 0, 0)
0092     RETAINED_FROM = datetime.datetime(2026, 7, 1)
0093
0094     def setUp(self) -> None:
0095         self.proxy = mock.MagicMock()
0096         self.proxy.get_db_utc_now.return_value = self.AS_OF
0097         self.proxy.get_partition_window.return_value = (5, self.RETAINED_FROM)
0098         self.proxy._commit.return_value = True
0099
0100     def test_rows(self) -> None:
0101         columns = DatasetSpec.columnNames().split(",")
0102         row = tuple(f"value of {column}" for column in columns)
0103         self.proxy.cur.fetchall.return_value = [row]
0104         success, message, data = MiscStandaloneModule.get_datasets_by_name(self.proxy, ["a_sub1", "b_sub2"], 1000)
0105         self.assertTrue(success, message)
0106         self.assertEqual(data, {"rows": [dict(zip(columns, row))], "as_of": self.AS_OF, "retained_from": self.RETAINED_FROM, "partitions": 5, "budget": 1000})
0107         self.assertIn("creationdate", data["rows"][0])
0108         sql, var_map = self.proxy.cur.execute.call_args[0]
0109         self.assertIn("FROM ATLAS_PANDA.Datasets WHERE name IN (", sql)
0110         self.assertEqual(sorted(var_map.values()), ["a_sub1", "b_sub2"])
0111
0112     def test_over_budget(self) -> None:
0113         success, message, data = MiscStandaloneModule.get_datasets_by_name(self.proxy, [f"a_sub{i}" for i in range(201)], 1000)
0114         self.assertFalse(success)
0115         self.assertEqual(data, {"rejected": "probe_budget", "keys": 201, "partitions": 5, "budget": 1000})
0116         self.proxy.cur.execute.assert_not_called()
0117
0118     def test_database_error(self) -> None:
0119         self.proxy.cur.execute.side_effect = Exception("ORA-00942: table or view does not exist")
0120         success, message, data = MiscStandaloneModule.get_datasets_by_name(self.proxy, ["a_sub1"], 1000)
0121         self.assertEqual((success, data), (False, {"error": "ORA-00942"}))
0122
0123
0124 class TestTaskDetailsSplitRule(unittest.TestCase):
0125     def _details(self, split_rule: str | None) -> dict[str, Any]:
0126         row: list[str | None] = [None] * len(JediTaskSpec.attributes)
0127         row[JediTaskSpec.attributes.index("splitRule")] = split_rule
0128         proxy = mock.MagicMock()
0129         proxy.cur.fetchone.side_effect = [tuple(row), (0, 0, 0, 0)]
0130         proxy.getClobObj.side_effect = [(True, [["template"]]), (True, [['{"cliParams": "prun"}']])]
0131         details: dict[str, Any] = TaskUtilsModule.get_task_details_json(proxy, 1)
0132         return details
0133
0134     def test_raw_kept(self) -> None:
0135         raw = "CI=1,MF=1"
0136         details = self._details(raw)
0137         self.assertEqual(details["split_rule_raw"], raw)
0138         self.assertEqual(details["splitRule"], task_split_rules.decode_split_rule(raw))
0139         self.assertNotEqual(details["splitRule"], raw)
0140
0141     def test_null(self) -> None:
0142         details = self._details(None)
0143         self.assertIsNone(details["split_rule_raw"])
0144         self.assertIsNone(details["splitRule"])
0145
0146
0147 if __name__ == "__main__":
0148     unittest.main()