summaryrefslogtreecommitdiff
path: root/test-distbuild-helper
blob: baf8ce8270d7982dc4145fa007870f2f880c2c5b (plain)
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
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
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()