diff options
Diffstat (limited to 'test.py')
| -rw-r--r-- | test.py | 25 |
1 files changed, 17 insertions, 8 deletions
@@ -7,22 +7,26 @@ from femtoqueue import FemtoQueue, FemtoTask from tempfile import mkdtemp original_time_fn = time.time + + def set_time_mock(ts: float): def mock_time_fn(): return ts time.time = mock_time_fn + def reset_time_mock(): time.time = original_time_fn + class TestFemtoQueue(unittest.TestCase): def test_basic(self): dir = mkdtemp() - q = FemtoQueue(data_dir = dir, node_id = "node1") + q = FemtoQueue(data_dir=dir, node_id="node1") # Add a JSON task - some_data = json.dumps({ "foo": "bar" }) + some_data = json.dumps({"foo": "bar"}) q.push(some_data.encode("utf-8")) # The task should be in the queue with the correct payload @@ -41,10 +45,10 @@ class TestFemtoQueue(unittest.TestCase): def test_aborted(self): dir = mkdtemp() - q = FemtoQueue(data_dir = dir, node_id = "node1") + q = FemtoQueue(data_dir=dir, node_id="node1") # Add a JSON task - some_data = json.dumps({ "foo": "bar" }) + some_data = json.dumps({"foo": "bar"}) q.push(some_data.encode("utf-8")) # The task should be in the queue with the correct payload @@ -55,7 +59,7 @@ class TestFemtoQueue(unittest.TestCase): self.assertEqual(parsed_data["foo"], "bar") # Simulate a fault and start over - q = FemtoQueue(data_dir = dir, node_id = "node1") + q = FemtoQueue(data_dir=dir, node_id="node1") # The task should still be assigned to this node task = q.pop() @@ -77,14 +81,18 @@ class TestFemtoQueue(unittest.TestCase): set_time_mock(0) # Node1 creates and claims the task - q1 = FemtoQueue(data_dir=dir, node_id="node1", timeout_stale_ms=timeout_stale_ms) + q1 = FemtoQueue( + data_dir=dir, node_id="node1", timeout_stale_ms=timeout_stale_ms + ) q1.push(b"stuck") task = q1.pop() self.assertIsNotNone(task) task = cast(FemtoTask, task) # Assert that node2 can not see the task - q2 = FemtoQueue(data_dir=dir, node_id="node2", timeout_stale_ms=timeout_stale_ms) + q2 = FemtoQueue( + data_dir=dir, node_id="node2", timeout_stale_ms=timeout_stale_ms + ) non_existant_task = q2.pop() self.assertIsNone(non_existant_task) @@ -162,5 +170,6 @@ class TestFemtoQueue(unittest.TestCase): # Queue should now be empty self.assertIsNone(q.pop()) -if __name__ == '__main__': + +if __name__ == "__main__": unittest.main() |