Once you’ve made sure that your message dispatching and application logic works
fine, you can actually start sending real messages to your process and checking
real replies. This can be done for single processes—I call this process
testing—and for your complete application (system testing).
When you test a single process, you create sockets that mimic all processes the
tested process communicates with. When you do a system test, you only mimic a
client or just invoke your program from the command line and check its output
(e.g., what it prints to stdtout and stderr or results written to a
database).
I’ll start with process testing, which is a bit more generalizable than system
testing.
Process Testing
The biggest problem I ran into when I started testing processes was that I
often made blocking calls to recv methods and these halted my tests and gave
me no output about what actually went wrong. Though you can make them non-
blocking by passing zmq.NOBLOCK as an extra argument, this doesn’t solve
your problems. You will now need a very precise timing and many
time.sleep(x) calls, because recv will instantly raise an error if there
is nothing to be received.
My solution for this was to wrap PyZMQ sockets and add a timeout to its send
and recv methods. The following wrapper will try to receive something for
one second and raise an exception if that failed. There’s also a simple
wrapper for methods like connect or bind,
but it’s really not that interesting, so I’ll omit it here.
# test/support.py
def get_wrapped_fwd(func):
"""
Returns a wrapper, that tries to call *func* multiple time in non-blocking
mode before rasing an :class:`zmq.ZMQError`.
"""
def forwarder(*args, **kwargs):
# 100 tries * 0.01 second == 1 second
for i in range(100):
try:
rep = func(*args, flags=zmq.NOBLOCK, **kwargs)
return rep
except zmq.ZMQError:
time.sleep(0.01)
# We should not get here, so raise an error.
msg = 'Could not %s message.' % func.__name__[:4]
raise zmq.ZMQError(msg)
return forwarder
This wrapper is now used to create a TestSocket class with the desired
behavior:
# test/support.py
class TestSocket(object):
"""
Wraps ZMQ :class:`~zmq.core.socket.Socket`. All *recv* and *send* methods
will be called multiple times in non-blocking mode before a
:class:`zmq.ZMQError` is raised.
"""
def __init__(self, context, sock_type):
self._context = context
sock = context.socket(sock_type)
self._sock = sock
forwards = [ # These methods can simply be forwarded
sock.bind,
sock.bind_to_random_port,
sock.connect,
sock.close,
sock.setsockopt,
]
wrapped_fwd = [ # These methods are wrapped with a for loop
sock.recv,
sock.recv_json,
sock.recv_multipart,
sock.recv_unicode,
sock.send,
sock.send_json,
sock.send_multipart,
sock.send_unicode,
]
for func in forwards:
setattr(self, func.__name__, get_forwarder(func))
for func in wrapped_fwd:
setattr(self, func.__name__, get_wrapped_fwd(func))
In order to reuse the same ports for all test methods, you need to cleanly
close all sockets after each test. To handle method level setup/teardown in
pytest, you need to implement a setup_method and a teardown_method. In the
setup method, you create one or more TestSocket instances that mimic other
processes and you also start the process to be tested:
# test/process/test_pongproc.py
import pytest
import zmq
from test.support import ProcessTest, make_sock
import pongproc
host = '127.0.0.1'
port = 5678
class TestProngProc(ProcessTest):
"""Communication test for the Platform Manager process."""
def setup_method(self, method):
"""
Creates and starts a PongProc process and sets up sockets to
communicate with it.
"""
self.context = zmq.Context()
# make_sock creates and connects a TestSocket that we will use to
# mimic the Ping process
self.req_sock = make_sock(self.context, zmq.REQ,
connect=(host, port))
self.pp = pongproc.PongProc((host, port))
self.pp.start()
def teardown_method(self, method):
"""
Sends a kill message to the pp and waits for the process to terminate.
"""
# Send a stop message to the prong process and wait until it joins
self.req_sock.send_multipart([b'["plzdiekthxbye", null]'])
self.pp.join()
self.req_sock.close()
You may have noticed that our test class inherits ProcessTests. This class
and some helpers in a conftest.py
allow us to use some magic that improves the readability of the actual test:
# test/process/test_pongproc.py
def test_ping(self):
"""Tests a ping-pong sequence."""
yield ('send', self.req_sock, [], ['ping', 1])
reply = yield ('recv', self.req_sock)
assert reply == [['pong', 1]]
You can just yield send or recv events from your test case! When you yield
a send, the test machinery tries to send a message via the specified socket.
When you yield a receive, ProcessTest tries to receive something from the
socket and sends its result back to your test function, so that you can easily
compare the reply with the expected result.
The example above is roughly equivalent to the following code:
self.req_sock.send_multipart([] + [json.dumps(['ping', 1])])
reply = self.req_sock.recv_multipart()
reply[-1] = json.loads[reply[-1]]
assert reply == [['pong', 1]]
So how does this work? By default, if pytests finds a test function that is a
generator, it assumes that it generates further test functions. Hence, our
first step is to override this behavior. We can do this in a conftest.py
file in the test/process/ directory by implementing a
pytest_pycollect_makeitem function. In this case, we collect generator
functions like normal functions:
# test/process/conftest.py
from inspect import isfunction, isgeneratorfunction
def pytest_pycollect_makeitem(collector, name, obj):
"""
Collects all instance methods that are generators and returns them as
normal function items.
"""
if collector.funcnamefilter(name) and hasattr(obj, '__call__'):
if isfunction(obj) or isgeneratorfunction(obj):
return collector._genfunctions(name, obj)
Now, we need to tell pytest how to run a test on the collected generator
functions. This can be done by implementing pytest_runtest_call. If the
object we are going to test (item.obj) is a generator function, we call the
run method of the object’s instance (item.obj.__self__.run) and pass the
generator function to it. If the test item contains a normal function, we run
the default test.
# test/process/conftest.py
def pytest_runtest_call(item):
"""
Passes the test generator (``item.obj``) to the ``run()`` method of the
generator's instance. This method should be inherited from
:class:`test.support.ProcessTest`.
"""
if isgeneratorfunction(item.obj):
item.obj.__self__.run(item.obj)
else: # Normal test execution for normal instance methods
item.runtest()
But wait—we didn’t implement a run method in our test case! So it must be
inherited from ProcessTest. Let’s take a look at it:
# test/support.py
class ProcessTest(object):
"""
Base class for process tests. It offers basic actions for sending and
receiving messages and implements the *run* methods that handles the
actual test generators.
"""
def run(self, testfunc):
"""
Iterates over the *testfunc* generator and executes all actions it
yields. Results will be sent back into the generator.
:param testfunc: A generator function that yields tuples containing
an action keyword, which should be a function of this or
the inheriting class (like ``send`` or ``recv``) and additional
parameters that will be passed to that function, e.g.:
``('send', socket_obj, ['header'], 'body')``
:type testfunc: generatorfunction
"""
item_gen = testfunc()
item = next(item_gen)
def throw_err(skip_levels=0):
"""
Throws the last error to *item_gen* and skips *skip_levels* in
the traceback to point to the line that yielded the last event.
"""
etype, evalue, tb = sys.exc_info()
for i in range(skip_levels):
tb = tb.tb_next
item_gen.throw(etype, evalue, tb)
try:
while True:
try:
# Call the event handler and pass the args,
# e.g., self.send(socket_obj, header, body)
ret = getattr(self, item[0])(*item[1:])
# Send the results back to the test and get the next item
item = item_gen.send(ret)
except zmq.ZMQError:
throw_err(3) # PyZMQ could not send/recv
except AssertionError:
throw_err(1) # Error in the test
except StopIteration:
pass
The run method simply iterates over all events our testfunc generates and
calls a method with the name of the event (e.g., send or recv). Their
return value is sent back into the generator. If an error occurs, the
exception’s traceback is modified to point to the line of code that yielded the
according event and not to the run method itself.
The methods send and recv roughly do the same as the snippet I showed you
above:
# test/support.py
def send(self, socket, header, body, extra_data=[]):
"""
JSON-encodes *body*, concatenates it with *header*, appends
*extra_data* and sends it as multipart message over *socket*.
*header* and *extra_data* should be lists containg byte objects or
objects implementing the buffer interface (like NumPy arrays).
"""
socket.send_multipart(header + [json.dumps(body)] + extra_data)
def recv(self, socket, json_load_index=-1):
"""
Receives and returns a multipart message from *socket* and tries to
JSON-decode the item at position *json_load_index* (defaults to ``-1``;
the last element in the list). The original byte string will be
replaced by the loaded object. Set *json_load_index* to ``None`` to get
the original, unchanged message.
"""
msg = socket.recv_multipart()
if json_load_index is not None:
msg[json_load_index] = json.loads(msg[json_load_index])
return msg
You can even add your own event handler to your test class. I used this, for
example, to add a log event that checks if a PyZMQ log handler sent the
expected log messages:
def log(self, substr=''):
"""
Receives a message and asserts, that it is a log message and that
*substr* is in that message.
Usage:
yield ('log', 'Ai iz in ur log mesage')
"""
msg = self.log_sock.recv_json()
assert msg[0] == 'log_message'
assert substr in msg[1]
What if your process starts further subprocesses?
In some cases, the process you are about to test starts additional subprocesses
that you don’t want to test. Even worse, these processes might communicate via
sockets bound to random ports. And EVEN WORSE, the process you are testing
might depend on excepting a KeyboardInterrupt to send stop messages to child
processes or to clean something up!
The last problem is quite easy to solve: You just a send a SIGINT to your
process from the test:
import os, signal
def teardown_method(self, method):
os.kill(self.my_proc.pid, signal.SIGINT)
self.my_proc.join()
# Now you can close the test sockets
If you don’t want to start a certain subprocess, you can just mock it.
Imagine, you have two processes a.A and b.B, where A starts B,
then you just mock B before starting A:
with mock.patch('b.B'):
self.a = A()
self.a.start()
Imagine now, that A binds a socket to a random port and uses that socket to
communicate with B. If you want to mock B in your tests, you need that port
number in order to connect to it and send messages to A.
But how can you get that number? When A creates B, it already runs in its
own process, so a simple attribute access won’t work. Setting a random seed
would only work if you did that directly in A when it’s already running. But
doing that just for the tests is not such a good idea. It also may not work
reliably on all systems and Python versions.
However, A must pass the socket number to B, so that B can connect to
A. Thus, we can create a mock for B that will send us its port number via a
queue <http://docs.python.org/py3k/library/multiprocessing#exchanging-objects-
between-processes>:
class ProcMock(mock.Mock):
"""
This mock returns itself when called, so it acts like both, the
process’ class and instance object.
"""
def __init__(self):
super().__init__()
self.queue = multiprocessing.Queue()
def __call__(self, port):
"""Will be called when A instantiates B and passes its port number."""
self.queue.put(port)
return self
def start(self):
return # Just make sure the methods exists and returns nothing
def join(self):
return # Just make sure the methods exists and returns nothing
class TestA(ProcessTest):
def setup_method(self):
b_mock = ProcMock()
with mock.patch('b.B', new=b_mock):
self.a = A()
self.a.start()
# Get the port A is listening on
port = b_mock.queue.get()
# ...
As you’ve seen, process testing is really not as simple as unit testing. But I
always found bugs with it that my unit tests coudn’t detect. If you cover all
communication sequences for a process in a process test, you can be pretty
sure, that it will also work flawlessly in the final application.
System Testing
If your application consists of more than one process, you still need to test
whether all processes work nicely together or not. This is something you cannot
simulate reliably with a process tests, as much as unit tests can’t replace the
process test.
Writing a good system test is very application-specific and can, depending on
the complexity of your application, be very hard or very easy. Fortunately, the
latter is the case for our ping-pong app. We just start it and copy its output
to a file. If the output is not what we expected, we modify the file
accordingly. In our test, we can now simply invoke our programm again, capture
its output and compare it to the contents of the file we created before:
# test/system/test_pongproc.py
import os.path
import subprocess
import pytest
def test_pongproc():
filename = os.path.join('test', 'data', 'pongproc.out')
expected = open(filename).read()
output = subprocess.check_output(['python', 'pongproc.py'],
universal_newlines=True)
assert output == expected
If your application was a server, another way of doing the system test would be
to emulate a client that speaks with it. Your system test would then be very
similar to your process tests, except that you only mimic the client and not
all processes your main process communicates with.
Other applications (like, for instance, simulations) might create a database
containing collected data. Here, you might check if these results match your
expectations.
Of course you can also combine these possibilities or do something completely
different …
