distbuild: Add test suite for distbuild-helper

This is mostly to check that the 'cancel entire subprocess tree' works
as expected. Revert that patch and the test fails.

There are also some tweaks included in this commit.

Change-Id: If297522e6589ebb3a07dac66a39eb243789e53aa

1 files changed, 205 insertions, 0 deletions

diff --git a/test-distbuild-helper b/test-distbuild-helper
new file mode 100755
index 00000000..baf8ce82
--- /dev/null
+++ b/test-distbuild-helper
@@ -0,0 +1,205 @@
+#!/usr/bin/python
+#
+# test-distbuild-helper.py -- tests for distbuild-helper tool
+#
+# Copyright (C) 2015  Codethink Limited
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation; version 2 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along
+# with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+
+import logging
+import os
+import subprocess
+import sys
+import unittest
+
+import distbuild
+
+
+def assert_process_is_running(process_id):
+    assert os.path.exists('/proc/%s' % process_id), \
+        'Expected process ID %s to be running, but it is not!' % process_id
+
+
+def assert_process_is_not_running(process_id):
+    assert not os.path.exists('/proc/%s' % process_id), \
+        'Expected process ID %s to not be running, but it is!' % process_id
+
+
+class DistbuildHelperTestCase(unittest.TestCase):
+    '''Base class for `distbuild-helper` test cases.'''
+
+    def run_distbuild_helper(self, loop):
+        '''Run `distbuild-helper`, and wait for it to connect to this process.
+
+        An unused port is allocated for communication with the
+        `distbuild-helper` process.
+
+        Returns a subprocess.Popen instance, and the port number that was
+        allocated.
+
+        '''
+
+        logging.info('Setting up a listener for connections from the '
+                     '`distbuild-helper` process.')
+        listener = distbuild.ListenServer(
+            addr='localhost', port=0, machine=distbuild.JsonRouter)
+        loop.add_state_machine(listener)
+
+        logging.info('Starting the `distbuild-helper` process.')
+        process = subprocess.Popen(
+            ['./distbuild-helper', '--parent-address', 'localhost',
+             '--parent-port', str(listener._port)])
+
+        logging.info('Waiting for connection from helper subprocess.')
+        json_router = loop.run_until_new_state_machine(
+            distbuild.JsonRouter)
+
+        return process, listener._port
+
+    def connect_to_helper(self, loop, port):
+        '''Returns a FakeWorker connected to the distbuild helper.
+
+        I hope this is not needed, and we can use the JsonMachine to talk to
+        the distbuild-helper process directly.'''
+
+        class ListenableJsonMachine(distbuild.JsonMachine):
+            '''Kludgy adapter around JsonMachine class.
+
+            The JsonMachine class is used for two-way communications over a
+            socket using JSON messages. This wrapper exists so we can have a
+            JsonMachine instance created by a ListenServer class. It simply
+            makes the prototype of __init__() match what ListenServer expects.
+
+            '''
+            def __init__(self, cm, conn):
+                super(ListenableJsonMachine, self).__init__(conn)
+
+        worker_connection_machine = distbuild.ConnectionMachine(
+            addr='localhost', port=port, machine=ListenableJsonMachine,
+            extra_args=[])
+        loop.add_state_machine(worker_connection_machine)
+
+        logging.info('Awaiting connection to worker port.')
+        jm = loop.run_until_new_state_machine(ListenableJsonMachine)
+
+        return jm
+
+    def assert_json_message(self, type):
+        '''Assert about the next message we receive from the helper.
+
+        Return the message that was received.
+
+        '''
+        event = self.helper_jm.mainloop.run_until_event(
+            self.helper_jm, distbuild.JsonNewMessage)
+        logging.debug('Received %s', event.msg)
+
+        assert event.msg['type'] == type, \
+            "Expected a JSON message of type %s from the helper, but " \
+            "received: %s" % (type, event.msg)
+
+        return event.msg
+
+    def send_exec_cancel(self, id):
+        '''Send an exec-cancel message to the helper process.'''
+        msg = distbuild.message('exec-cancel', id=id)
+        self.helper_jm.send(msg)
+
+    def send_exec_request(self, id, argv):
+        '''Send an exec-request message to the helper process.'''
+        msg = distbuild.message(
+            'exec-request', id=id, argv=argv, stdin_contents='',
+        )
+        self.helper_jm.send(msg)
+
+
+class ExecutionTests(DistbuildHelperTestCase):
+    '''Test the execution and cancellation of subprocesses.'''
+
+    def setUp(self):
+        #logging.basicConfig(level=logging.DEBUG, stream=sys.stdout)
+
+        loop = distbuild.mainloop.TestableMainLoop()
+        self.helper_process, helper_port = self.run_distbuild_helper(loop)
+        self.helper_jm = self.connect_to_helper(loop, helper_port)
+
+    def tearDown(self):
+        self.helper_process.terminate()
+
+    def test_run_process_success(self):
+        '''Test that exec-request starts a process, and we get a response.'''
+
+        TEST_ARGV = ['sh', '-c', "echo test"]
+
+        self.send_exec_request(id='test1', argv=TEST_ARGV)
+
+        msg = self.assert_json_message(type='exec-output')
+        self.assertEqual(msg['stdout'], 'test\n')
+
+        msg = self.assert_json_message(type='exec-response')
+        self.assertEqual(msg['exit'], 0)
+
+    def test_cancel_process_tree(self):
+        '''Test that exec-cancel kills the entire process tree.
+
+        By default SIGKILL only kills the direct child process, but this may
+        have spawned other processes that might keep running.
+
+        To test that the process tree is killed, this test runs a script which
+        spawns another shell. The 2nd shell will output its process ID, then
+        sleep for 10 seconds, then output another message. If we see the second
+        message then we know that the SIGKILL didn't work correctly -- if the
+        2nd shell had been killed, it would not have been able to output the
+        second message.
+
+        '''
+
+        # This shell script runs a subprocess that will sleep for 10 seconds,
+        # then output a line of text, unless killed during the sleep.
+
+        TEST_PROGRAM = '''
+        sh -c 'echo "$$ is child process ID";
+               sleep 10;
+               echo "Process was not killed."'
+        '''
+
+        # Start the subprocess and wait for the first line of output.
+
+        self.send_exec_request(id='test1', argv=['sh', '-c', TEST_PROGRAM])
+
+        msg = self.assert_json_message(type='exec-output')
+
+        stdout = msg['stdout']
+        self.assertIn('child process ID', stdout)
+
+        child_process_id = stdout.split()[0]
+        assert_process_is_running(child_process_id)
+
+        # Now cancel the subprocess, while it is in the 'sleep' command, and
+        # wait for the exec-response message. If we receive an exec-output
+        # message here instead then we know that the process wasn't killed
+        # during the 'sleep'.
+
+        self.send_exec_cancel(id='test1')
+
+        msg = self.assert_json_message(type='exec-response')
+
+        assert_process_is_not_running(child_process_id)
+
+        assert msg['exit'] == -9, \
+            'Process was not killed -- exit code is %i' % msg['exit']
+
+
+if __name__ == '__main__':
+    unittest.main()