Skip to content

Commit 40b76d7

Browse files
mhartmaymhartmay
committed
first_stage_test: Add more tests
 + test_stdin_non_blocking + test_stdin_blocking + test_stage Signed-off-by: Marc Hartmayer <[email protected]>
1 parent fff2507 commit 40b76d7

File tree

1 file changed

+199
-0
lines changed

1 file changed

+199
-0
lines changed

‎tests/first_stage_test.py‎

Lines changed: 199 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,6 @@
1+
import fcntl
2+
import os
3+
import select
14
import subprocess
25

36
try:
@@ -93,3 +96,199 @@ def test_valid_syntax(self):
9396
)
9497
finally:
9598
fp.close()
99+
100+
def test_stage(self):
101+
"""Test that first stage works
102+
103+
The boot command should read the preamble from STDIN, write all ECO
104+
markers to STDOUT, and then execute the preamble.
105+
106+
This test writes the complete preamble to STDIN.
107+
108+
1. Fork child reads from STDIN
109+
2. Fork child writes all ECO markers to stdout as expected.
110+
111+
"""
112+
113+
proc = subprocess.Popen(
114+
args=self.args,
115+
stdout=subprocess.PIPE,
116+
stderr=subprocess.PIPE,
117+
stdin=subprocess.PIPE,
118+
)
119+
try:
120+
try:
121+
stdout, stderr = proc.communicate(input=self.preamble, timeout=10)
122+
except TypeError:
123+
stdout, stderr = proc.communicate(input=self.preamble)
124+
except:
125+
proc.kill()
126+
self.fail("First stage did not finish")
127+
else:
128+
# proc is killed by SIGTERM as the broker kills itself
129+
self.assertEqual(-15, proc.returncode)
130+
self.assertTrue(
131+
stdout.startswith(
132+
mitogen.parent.BootstrapProtocol.EC0_MARKER
133+
+ b("\n")
134+
+ mitogen.parent.BootstrapProtocol.EC1_MARKER
135+
+ b("\n")
136+
+ mitogen.parent.BootstrapProtocol.EC2_MARKER
137+
+ b("\n"),
138+
)
139+
)
140+
self.assertIn(b'shutting down', stdout)
141+
self.assertEqual(
142+
b(""),
143+
stderr,
144+
)
145+
146+
def test_stdin_non_blocking(self):
147+
"""Test that first stage works with non-blocking STDIN
148+
149+
The boot command should read the preamble from STDIN, write all ECO
150+
markers to STDOUT, and then execute the preamble.
151+
152+
This test writes the complete preamble to non-blocking STDIN.
153+
154+
1. Fork child reads from non-blocking STDIN
155+
2. Fork child writes all ECO markers to stdout as expected.
156+
157+
"""
158+
159+
PIPE_SIZE = 4096
160+
# Make sure that not all data can be written in one write operation.
161+
CHUNK_SIZE = 2 * PIPE_SIZE
162+
r, w = mitogen.core.pipe(blocking=False)
163+
with w:
164+
try:
165+
fcntl.fcntl(w.fileno(), fcntl.F_SETPIPE_SZ, PIPE_SIZE)
166+
except AttributeError:
167+
pass
168+
else:
169+
self.assertEqual(fcntl.fcntl(w.fileno(), fcntl.F_GETPIPE_SZ), PIPE_SIZE)
170+
171+
proc = subprocess.Popen(
172+
args=self.args,
173+
stdout=subprocess.PIPE,
174+
stderr=subprocess.PIPE,
175+
stdin=r,
176+
close_fds=True,
177+
)
178+
# Close the read fd in the parent
179+
r.close()
180+
181+
view = mitogen.core.BufferType(self.preamble, 0)
182+
offset = 0
183+
while offset < len(view):
184+
end = min(offset + CHUNK_SIZE, len(view))
185+
_, w_fds, _ = select.select([], [w.fileno()], [], 10)
186+
self.assertIn(w.fileno(), w_fds)
187+
written = os.write(w.fileno(), view[offset:end])
188+
self.assertGreater(written, 0)
189+
offset += written
190+
191+
try:
192+
try:
193+
returncode = proc.wait(timeout=10)
194+
except TypeError:
195+
returncode = proc.wait()
196+
except:
197+
proc.kill()
198+
self.fail("First stage did not finish")
199+
else:
200+
try:
201+
stdout = proc.stdout.read()
202+
# proc is killed by SIGTERM as the broker kills itself
203+
self.assertEqual(-15, proc.returncode)
204+
self.assertTrue(
205+
stdout.startswith(
206+
mitogen.parent.BootstrapProtocol.EC0_MARKER
207+
+ b("\n")
208+
+ mitogen.parent.BootstrapProtocol.EC1_MARKER
209+
+ b("\n")
210+
+ mitogen.parent.BootstrapProtocol.EC2_MARKER
211+
+ b("\n"),
212+
)
213+
)
214+
self.assertIn(b'shutting down', stdout)
215+
self.assertEqual(
216+
b(""),
217+
proc.stderr.read(),
218+
)
219+
finally:
220+
proc.stdout.close()
221+
proc.stderr.close()
222+
223+
def test_stdin_blocking(self):
224+
"""Test that first stage works with blocking STDIN
225+
226+
The boot command should read the preamble from STDIN, write all ECO
227+
markers to STDOUT, and then execute the preamble.
228+
229+
This test writes the complete preamble to blocking STDIN.
230+
231+
1. Fork child reads from blocking STDIN
232+
2. Fork child writes all ECO markers to stdout as expected.
233+
234+
"""
235+
PIPE_SIZE = 4096
236+
# Make sure that not all data can be written in one write operation.
237+
CHUNK_SIZE = 2 * PIPE_SIZE
238+
r, w = mitogen.core.pipe(blocking=True)
239+
with w:
240+
try:
241+
fcntl.fcntl(w.fileno(), fcntl.F_SETPIPE_SZ, PIPE_SIZE)
242+
except AttributeError:
243+
pass
244+
else:
245+
self.assertEqual(fcntl.fcntl(w.fileno(), fcntl.F_GETPIPE_SZ), PIPE_SIZE)
246+
proc = subprocess.Popen(
247+
args=self.args,
248+
stdout=subprocess.PIPE,
249+
stderr=subprocess.PIPE,
250+
stdin=r,
251+
close_fds=True,
252+
)
253+
# Close the read fd in the parent
254+
r.close()
255+
256+
view = mitogen.core.BufferType(self.preamble, 0)
257+
offset = 0
258+
while offset < len(view):
259+
end = min(offset + CHUNK_SIZE, len(view))
260+
written = os.write(w.fileno(), view[offset:end])
261+
self.assertGreater(written, 0)
262+
offset += written
263+
264+
try:
265+
try:
266+
returncode = proc.wait(timeout=10)
267+
except TypeError:
268+
returncode = proc.wait()
269+
except:
270+
proc.kill()
271+
self.fail("First stage did not finish")
272+
else:
273+
try:
274+
stdout = proc.stdout.read()
275+
# proc is killed by SIGTERM as the broker kills itself
276+
self.assertEqual(-15, proc.returncode)
277+
self.assertTrue(
278+
stdout.startswith(
279+
mitogen.parent.BootstrapProtocol.EC0_MARKER
280+
+ b("\n")
281+
+ mitogen.parent.BootstrapProtocol.EC1_MARKER
282+
+ b("\n")
283+
+ mitogen.parent.BootstrapProtocol.EC2_MARKER
284+
+ b("\n"),
285+
)
286+
)
287+
self.assertIn(b'shutting down', stdout)
288+
self.assertEqual(
289+
b(""),
290+
proc.stderr.read(),
291+
)
292+
finally:
293+
proc.stdout.close()
294+
proc.stderr.close()

0 commit comments

Comments
 (0)