1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
|
from typing import cast
import os
import time
import unittest
import json
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")
# Add a JSON task
some_data = json.dumps({ "foo": "bar" })
q.push(some_data.encode("utf-8"))
# The task should be in the queue with the correct payload
task = q.pop()
self.assertIsNotNone(task)
task = cast(FemtoTask, task)
parsed_data = json.loads(task.data.decode("utf-8"))
self.assertEqual(parsed_data["foo"], "bar")
# Mark the task as done
q.done(task)
# There should be no tasks available
task = q.pop()
self.assertIsNone(task)
def test_aborted(self):
dir = mkdtemp()
q = FemtoQueue(data_dir = dir, node_id = "node1")
# Add a JSON task
some_data = json.dumps({ "foo": "bar" })
q.push(some_data.encode("utf-8"))
# The task should be in the queue with the correct payload
task = q.pop()
self.assertIsNotNone(task)
task = cast(FemtoTask, task)
parsed_data = json.loads(task.data.decode("utf-8"))
self.assertEqual(parsed_data["foo"], "bar")
# Simulate a fault and start over
q = FemtoQueue(data_dir = dir, node_id = "node1")
# The task should still be assigned to this node
task = q.pop()
self.assertIsNotNone(task)
task = cast(FemtoTask, task)
parsed_data = json.loads(task.data.decode("utf-8"))
self.assertEqual(parsed_data["foo"], "bar")
# Mark the task as done
q.done(task)
# There should be no tasks available
task = q.pop()
self.assertIsNone(task)
def test_release_stale_tasks(self):
dir = mkdtemp()
timeout_stale_ms = 100
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.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)
non_existant_task = q2.pop()
self.assertIsNone(non_existant_task)
# Simulate the task becoming stale by skipping forward in time
set_time_mock(10)
# Now node2 should see the task as stale and reclaim it
revived_task = q2.pop()
self.assertIsNotNone(revived_task)
revived_task = cast(FemtoTask, revived_task)
self.assertEqual(revived_task.data, b"stuck")
reset_time_mock()
def test_mark_task_failed(self):
dir = mkdtemp()
q = FemtoQueue(data_dir=dir, node_id="node1")
# Push and pop a task
q.push(b"will fail")
task = q.pop()
self.assertIsNotNone(task)
task = cast(FemtoTask, task)
# Mark the task as failed
q.fail(task)
# Ensure the file now exists in the failed directory
failed_path = os.path.join(dir, "failed", task.id)
self.assertTrue(os.path.exists(failed_path))
# The task should not reappear in pop
self.assertIsNone(q.pop())
def test_mark_task_done(self):
dir = mkdtemp()
q = FemtoQueue(data_dir=dir, node_id="node1")
# Push and pop a task
q.push(b"complete me")
task = q.pop()
self.assertIsNotNone(task)
task = cast(FemtoTask, task)
# Mark the task as done
q.done(task)
# Ensure the file is now in the 'done' directory
done_path = os.path.join(dir, "done", task.id)
self.assertTrue(os.path.exists(done_path))
# The task should not be returned again
self.assertIsNone(q.pop())
def test_fifo_order(self):
dir = mkdtemp()
q = FemtoQueue(data_dir=dir, node_id="node1")
num_tasks = 100
# Push tasks with sequential numbers as content
for i in range(num_tasks):
q.push(str(i).encode("utf-8"))
last_value = -1
for _ in range(num_tasks):
task = q.pop()
self.assertIsNotNone(task)
task = cast(FemtoTask, task)
current_value = int(task.data.decode("utf-8"))
self.assertEqual(current_value, last_value + 1)
last_value = current_value
q.done(task)
# Queue should now be empty
self.assertIsNone(q.pop())
if __name__ == '__main__':
unittest.main()
|
