Skip to content

Commit b8ecfc8

Browse files
yokomotodyokomotod
authored
unit test for master.build (#26)
* add unit test for master.build() * remove unused * mv test_master_build to integration_test/
1 parent 4dab38b commit b8ecfc8

File tree

4 files changed

+171
-4
lines changed

4 files changed

+171
-4
lines changed

‎integration_test/__init__.py‎

Whitespace-only changes.

‎integration_test/test_master_build.py‎

Lines changed: 167 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,167 @@
1+
from __future__ import annotations
2+
3+
import time
4+
import unittest
5+
from unittest.mock import MagicMock
6+
7+
import gokart
8+
import luigi
9+
from kubernetes import client
10+
11+
from kannon import Kannon, TaskOnBullet
12+
13+
14+
class MockTaskOnKart(gokart.TaskOnKart):
15+
wait_sec = 1.
16+
started_at: float | None = None
17+
18+
def run(self) -> None:
19+
self.started_at = time.time()
20+
21+
def complete(self) -> bool:
22+
if self.started_at is None:
23+
return False
24+
25+
return time.time() > self.started_at + self.wait_sec
26+
27+
28+
class MockTaskOnBullet(TaskOnBullet):
29+
wait_sec = 1.
30+
started_at: float | None = None
31+
32+
def run(self) -> None:
33+
self.started_at = time.time()
34+
35+
def complete(self) -> bool:
36+
if self.started_at is None:
37+
return False
38+
39+
return time.time() > self.started_at + self.wait_sec
40+
41+
42+
class MockKannon(Kannon):
43+
44+
def __init__(self) -> None:
45+
super().__init__(
46+
api_instance=None,
47+
template_job=client.V1Job(metadata=client.V1ObjectMeta()),
48+
job_prefix="",
49+
path_child_script=__file__, # just pass any existing file as dummy
50+
env_to_inherit=None,
51+
)
52+
53+
def _exec_gokart_task(self, task: MockTaskOnKart) -> None:
54+
task.run()
55+
56+
def _exec_bullet_task(self, task: MockTaskOnBullet) -> None:
57+
task.run()
58+
59+
60+
class TestConsumeTaskQueue(unittest.TestCase):
61+
62+
def test_single_task_on_kart(self) -> None:
63+
self.maxDiff = None
64+
65+
class Example(MockTaskOnKart):
66+
pass
67+
68+
root_task = Example()
69+
70+
master = MockKannon()
71+
with self.assertLogs() as cm:
72+
master.build(root_task)
73+
74+
root_task_info = master._gen_task_info(root_task)
75+
self.assertEqual(cm.output, [
76+
'INFO:kannon.master:Creating task queue...',
77+
f'INFO:kannon.master:Task {root_task_info} is pushed to task queue',
78+
'INFO:kannon.master:Consuming task queue...',
79+
f'INFO:kannon.master:Checking if task {root_task_info} is executable...',
80+
f'INFO:kannon.master:Executing task {root_task_info} on master job...',
81+
f'INFO:kannon.master:Completed task {root_task_info} on master job.',
82+
'INFO:kannon.master:All tasks completed!',
83+
])
84+
85+
def test_single_task_on_bullet(self) -> None:
86+
self.maxDiff = None
87+
88+
class Example(MockTaskOnBullet):
89+
pass
90+
91+
root_task = Example()
92+
# FIXME: additional patch due to no sleep in this case
93+
root_task.complete = MagicMock(side_effect=[False, False, True]) # type:ignore
94+
95+
master = MockKannon()
96+
with self.assertLogs() as cm:
97+
master.build(root_task)
98+
99+
root_task_info = master._gen_task_info(root_task)
100+
self.assertEqual(cm.output, [
101+
'INFO:kannon.master:Creating task queue...',
102+
f'INFO:kannon.master:Task {root_task_info} is pushed to task queue',
103+
'INFO:kannon.master:Consuming task queue...',
104+
f'INFO:kannon.master:Checking if task {root_task_info} is executable...',
105+
f'INFO:kannon.master:Trying to run task {root_task_info} on child job...',
106+
f'INFO:kannon.master:Task {root_task_info} is still running on child job.',
107+
f'INFO:kannon.master:Task {root_task_info} is already completed.',
108+
'INFO:kannon.master:All tasks completed!',
109+
])
110+
111+
def test_three_task_on_bullet(self) -> None:
112+
self.maxDiff = None
113+
114+
class Child(MockTaskOnBullet):
115+
param = luigi.IntParameter()
116+
117+
c1 = Child(param=1)
118+
c1.wait_sec = 4
119+
c2 = Child(param=2)
120+
c2.wait_sec = 3
121+
c3 = Child(param=3)
122+
c3.wait_sec = 2
123+
124+
class Parent(MockTaskOnKart):
125+
126+
def requires(self) -> list[Child]:
127+
return [c1, c2, c3]
128+
129+
root_task = Parent()
130+
131+
master = MockKannon()
132+
with self.assertLogs() as cm:
133+
master.build(root_task)
134+
135+
c1_task_info = master._gen_task_info(c1)
136+
c2_task_info = master._gen_task_info(c2)
137+
c3_task_info = master._gen_task_info(c3)
138+
root_task_info = master._gen_task_info(root_task)
139+
self.assertEqual(cm.output, [
140+
'INFO:kannon.master:Creating task queue...',
141+
f'INFO:kannon.master:Task {c1_task_info} is pushed to task queue',
142+
f'INFO:kannon.master:Task {c2_task_info} is pushed to task queue',
143+
f'INFO:kannon.master:Task {c3_task_info} is pushed to task queue',
144+
f'INFO:kannon.master:Task {root_task_info} is pushed to task queue',
145+
'INFO:kannon.master:Consuming task queue...',
146+
f'INFO:kannon.master:Checking if task {c1_task_info} is executable...',
147+
f'INFO:kannon.master:Trying to run task {c1_task_info} on child job...',
148+
f'INFO:kannon.master:Checking if task {c2_task_info} is executable...',
149+
f'INFO:kannon.master:Trying to run task {c2_task_info} on child job...',
150+
f'INFO:kannon.master:Checking if task {c3_task_info} is executable...',
151+
f'INFO:kannon.master:Trying to run task {c3_task_info} on child job...',
152+
f'INFO:kannon.master:Checking if task {root_task_info} is executable...',
153+
f'INFO:kannon.master:Task {c1_task_info} is still running on child job.',
154+
f'INFO:kannon.master:Task {c2_task_info} is still running on child job.',
155+
f'INFO:kannon.master:Task {c3_task_info} is still running on child job.',
156+
f'INFO:kannon.master:Checking if task {root_task_info} is executable...',
157+
f'INFO:kannon.master:Executing task {root_task_info} on master job...',
158+
f'INFO:kannon.master:Completed task {root_task_info} on master job.',
159+
f'INFO:kannon.master:Task {c1_task_info} is already completed.',
160+
f'INFO:kannon.master:Task {c2_task_info} is already completed.',
161+
f'INFO:kannon.master:Task {c3_task_info} is already completed.',
162+
'INFO:kannon.master:All tasks completed!',
163+
])
164+
165+
166+
if __name__ == '__main__':
167+
unittest.main()

‎test/test_master.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
from kannon import Kannon
1010

1111

12-
class TestStringMethods(unittest.TestCase):
12+
class TestCreateTaskQueue(unittest.TestCase):
1313

1414
def test_create_task_queue(self) -> None:
1515

‎tox.ini‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,17 +15,17 @@ commands = yapf -dr . {posargs}
1515
[testenv:isort]
1616
allowlist_externals = isort
1717
skip_install = true
18-
commands = isort -c ./kannon ./test ./example {posargs}
18+
commands = isort -c ./kannon ./test ./integration_test ./example {posargs}
1919

2020
[testenv:flake8]
2121
allowlist_externals = pflake8
2222
skip_install = true
23-
commands = pflake8 ./kannon ./test ./example {posargs}
23+
commands = pflake8 ./kannon ./test ./integration_test ./example {posargs}
2424

2525
[testenv:mypy]
2626
allowlist_externals = mypy
2727
skip_install = true
28-
commands = mypy ./kannon ./test ./example {posargs}
28+
commands = mypy ./kannon ./test ./integration_test ./example {posargs}
2929

3030
[gh-actions]
3131
python =

0 commit comments

Comments
 (0)