tests/cas/test_client.py

@@ -26,7 +26,8 @@ from buildgrid.client.cas import download, upload
 from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
 from buildgrid.utils import create_digest
 
-from ..utils.cas import serve_cas, run_in_subprocess
+from ..utils.cas import serve_cas
+from ..utils.utils import run_in_subprocess
 
 
 INTANCES = ['', 'instance']

tests/cas/test_storage.py

@@ -32,7 +32,8 @@ from buildgrid.server.cas.storage.s3 import S3Storage
 from buildgrid.server.cas.storage.with_cache import WithCacheStorage
 from buildgrid.settings import HASH
 
-from ..utils.cas import serve_cas, run_in_subprocess
+from ..utils.cas import serve_cas
+from ..utils.utils import run_in_subprocess
 
 
 BLOBS = [(b'abc', b'defg', b'hijk', b'')]

tests/integration/bot_session.py

@@ -14,56 +14,174 @@
 
 # pylint: disable=redefined-outer-name
 
-import uuid
+import asyncio
 
+import grpc
 import pytest
 
-from buildgrid.bot import bot_session
+from buildgrid._enums import LeaseState
+from buildgrid._protos.google.devtools.remoteworkers.v1test2 import bots_pb2
+from buildgrid.bot.hardware.worker import Worker
+from buildgrid.bot.hardware.interface import HardwareInterface
+from buildgrid.bot.session import BotSession
+from buildgrid.bot.interface import BotInterface
 
+from ..utils.utils import run_in_subprocess
+from ..utils.bots_interface import serve_bots_interface
 
-@pytest.mark.parametrize("docker_value", ["True", "False"])
-@pytest.mark.parametrize("os_value", ["nexus7", "nexus8"])
-def test_create_device(docker_value, os_value):
-    properties = {'docker': docker_value, 'os': os_value}
-    device = bot_session.Device(properties)
 
-    assert uuid.UUID(device.name, version=4)
-    assert properties == device.properties
+INSTANCES = ['', 'instance']
 
 
-def test_create_device_key_fail():
-    properties = {'voight': 'kampff'}
+# Use subprocess to avoid creation of gRPC threads in main process
+# See https://github.com/grpc/grpc/blob/master/doc/fork_support.md
+# Multiprocessing uses pickle which protobufs don't work with
+# Workaround wrapper to send messages as strings
+class ServerInterface:
 
-    with pytest.raises(KeyError):
-        bot_session.Device(properties)
+    def __init__(self, remote):
+        self.__remote = remote
 
+    def create_bot_session(self, parent, bot_session):
 
-def test_create_device_value_fail():
-    properties = {'docker': True}
+        def __create_bot_session(queue, remote, parent, string_bot_session):
+            bot_session = bots_pb2.BotSession()
+            bot_session.ParseFromString(string_bot_session)
 
-    with pytest.raises(ValueError):
-        bot_session.Device(properties)
+            interface = BotInterface(grpc.insecure_channel(remote))
 
+            result = interface.create_bot_session(parent, bot_session)
+            queue.put(result.SerializeToString())
 
-def test_create_worker():
-    properties = {'pool': 'swim'}
-    configs = {'DockerImage': 'Windows'}
-    worker = bot_session.Worker(properties, configs)
+        string_bot_session = bot_session.SerializeToString()
+        result = run_in_subprocess(__create_bot_session,
+                                   self.__remote, parent, string_bot_session)
 
-    assert properties == worker.properties
-    assert configs == worker.configs
+        bot_session = bots_pb2.BotSession()
+        bot_session.ParseFromString(result)
+        return bot_session
 
-    device = bot_session.Device()
-    worker.add_device(device)
+    def update_bot_session(self, bot_session, update_mask=None):
 
-    assert worker._devices[0] == device
+        def __update_bot_session(queue, remote, string_bot_session, update_mask):
+            bot_session = bots_pb2.BotSession()
+            bot_session.ParseFromString(string_bot_session)
 
+            interface = BotInterface(grpc.insecure_channel(remote))
 
-def test_create_worker_key_fail():
-    properties = {'voight': 'kampff'}
-    configs = {'voight': 'kampff'}
+            result = interface.update_bot_session(bot_session, update_mask)
+            queue.put(result.SerializeToString())
 
-    with pytest.raises(KeyError):
-        bot_session.Worker(properties)
-    with pytest.raises(KeyError):
-        bot_session.Worker(configs)
+        string_bot_session = bot_session.SerializeToString()
+        result = run_in_subprocess(__update_bot_session,
+                                   self.__remote, string_bot_session, update_mask)
+
+        bot_session = bots_pb2.BotSession()
+        bot_session.ParseFromString(result)
+        return bot_session
+
+
+@pytest.mark.parametrize('instance', INSTANCES)
+def test_create_bot_session(instance):
+
+    with serve_bots_interface([instance]) as server:
+        interface = ServerInterface(server.remote)
+        hardware_interface = HardwareInterface(Worker())
+        session = BotSession(instance, interface, hardware_interface, None)
+        session.create_bot_session()
+        assert session.get_pb2() == server.get_bot_session()
+
+
+@pytest.mark.parametrize('instance', INSTANCES)
+def test_update_bot_session(instance):
+
+    with serve_bots_interface([instance]) as server:
+        interface = ServerInterface(server.remote)
+        hardware_interface = HardwareInterface(Worker())
+        session = BotSession(instance, interface, hardware_interface, None)
+        session.create_bot_session()
+        assert session.get_pb2() == server.get_bot_session()
+        session.update_bot_session()
+        assert session.get_pb2() == server.get_bot_session()
+
+
+@pytest.mark.parametrize('instance', INSTANCES)
+def test_create_bot_session_with_work(instance):
+
+    def __work(lease, context, event):
+        return lease
+
+    with serve_bots_interface([instance]) as server:
+        interface = ServerInterface(server.remote)
+        hardware_interface = HardwareInterface(Worker())
+        session = BotSession(instance, interface, hardware_interface, __work)
+        server.inject_work()
+        session.create_bot_session()
+
+        assert len(session.get_pb2().leases) == 1
+
+        loop = asyncio.get_event_loop()
+        for task in asyncio.Task.all_tasks():
+            loop.run_until_complete(task)
+
+        assert session.get_pb2().leases[0].state == LeaseState.COMPLETED.value
+
+
+@pytest.mark.parametrize('instance', INSTANCES)
+def test_update_bot_session_with_work(instance):
+
+    def __work(lease, context, event):
+        return lease
+
+    with serve_bots_interface([instance]) as server:
+        interface = ServerInterface(server.remote)
+        hardware_interface = HardwareInterface(Worker())
+        session = BotSession(instance, interface, hardware_interface, __work)
+        session.create_bot_session()
+        server.inject_work()
+        session.update_bot_session()
+
+        assert len(session.get_pb2().leases) == 1
+
+        loop = asyncio.get_event_loop()
+        for task in asyncio.Task.all_tasks():
+            loop.run_until_complete(task)
+
+        assert session.get_pb2().leases[0].state == LeaseState.COMPLETED.value
+
+
+@pytest.mark.parametrize('instance', INSTANCES)
+def test_cancel_leases(instance):
+
+    def __work(lease, context, cancel_event):
+        # while not cancel_event.is_set():
+
+        return lease
+
+    with serve_bots_interface([instance]) as server:
+        interface = ServerInterface(server.remote)
+        hardware_interface = HardwareInterface(Worker())
+        session = BotSession(instance, interface, hardware_interface, __work)
+
+        lease = bots_pb2.Lease()
+        lease.state = LeaseState.PENDING.value
+        lease.id = 'foo'
+        server.inject_work(lease)
+        session.create_bot_session()
+
+        leases_pb2 = session.get_pb2().leases
+        assert len(leases_pb2) == 1
+        assert leases_pb2[0].state == LeaseState.ACTIVE.value
+
+        server.cancel_lease(leases_pb2[0].id)
+        session.update_bot_session()
+        assert len(session.get_pb2().leases) == 1
+
+        loop = asyncio.get_event_loop()
+        for task in asyncio.Task.all_tasks():
+            try:
+                loop.run_until_complete(task)
+            except asyncio.CancelledError:
+                pass
+
+        assert session.get_pb2().leases[0].state == LeaseState.CANCELLED.value

tests/integration/bots_service.py

@@ -17,7 +17,6 @@
 
 # pylint: disable=redefined-outer-name
 
-import copy
 from unittest import mock
 
 import grpc
@@ -150,129 +149,6 @@ def test_update_leases_with_work(bot_session, context, instance):
     assert response_action == action_digest
 
 
-def test_update_leases_work_complete(bot_session, context, instance):
-    request = bots_pb2.CreateBotSessionRequest(parent='',
-                                               bot_session=bot_session)
-    # Create bot session
-    # Simulated the severed binding between client and server
-    response = copy.deepcopy(instance.CreateBotSession(request, context))
-
-    # Inject work
-    action_digest = remote_execution_pb2.Digest(hash='gaff')
-    _inject_work(instance._instances[""]._scheduler, action_digest=action_digest)
-
-    request = bots_pb2.UpdateBotSessionRequest(name=response.name,
-                                               bot_session=response)
-    response = copy.deepcopy(instance.UpdateBotSession(request, context))
-
-    assert response.leases[0].state == LeaseState.PENDING.value
-    response.leases[0].state = LeaseState.ACTIVE.value
-
-    request = bots_pb2.UpdateBotSessionRequest(name=response.name,
-                                               bot_session=response)
-
-    response = copy.deepcopy(instance.UpdateBotSession(request, context))
-
-    response.leases[0].state = LeaseState.COMPLETED.value
-    response.leases[0].result.Pack(remote_execution_pb2.ActionResult())
-
-    request = bots_pb2.UpdateBotSessionRequest(name=response.name,
-                                               bot_session=response)
-    response = copy.deepcopy(instance.UpdateBotSession(request, context))
-
-    assert len(response.leases) is 0
-
-
-def test_work_rejected_by_bot(bot_session, context, instance):
-    request = bots_pb2.CreateBotSessionRequest(parent='',
-                                               bot_session=bot_session)
-    # Inject work
-    action_digest = remote_execution_pb2.Digest(hash='gaff')
-    _inject_work(instance._instances[""]._scheduler, action_digest=action_digest)
-
-    # Simulated the severed binding between client and server
-    response = copy.deepcopy(instance.CreateBotSession(request, context))
-
-    # Reject work
-    assert response.leases[0].state == LeaseState.PENDING.value
-    response.leases[0].state = LeaseState.COMPLETED.value
-    request = bots_pb2.UpdateBotSessionRequest(name=response.name,
-                                               bot_session=response)
-
-    response = instance.UpdateBotSession(request, context)
-
-    context.set_code.assert_called_once_with(grpc.StatusCode.UNIMPLEMENTED)
-
-
-@pytest.mark.parametrize("state", [LeaseState.LEASE_STATE_UNSPECIFIED, LeaseState.PENDING])
-def test_work_out_of_sync_from_pending(state, bot_session, context, instance):
-    request = bots_pb2.CreateBotSessionRequest(parent='',
-                                               bot_session=bot_session)
-    # Inject work
-    action_digest = remote_execution_pb2.Digest(hash='gaff')
-    _inject_work(instance._instances[""]._scheduler, action_digest=action_digest)
-
-    # Simulated the severed binding between client and server
-    response = copy.deepcopy(instance.CreateBotSession(request, context))
-
-    response.leases[0].state = state.value
-
-    request = bots_pb2.UpdateBotSessionRequest(name=response.name,
-                                               bot_session=response)
-
-    response = instance.UpdateBotSession(request, context)
-
-    context.set_code.assert_called_once_with(grpc.StatusCode.DATA_LOSS)
-
-
-@pytest.mark.parametrize("state", [LeaseState.LEASE_STATE_UNSPECIFIED, LeaseState.PENDING])
-def test_work_out_of_sync_from_active(state, bot_session, context, instance):
-    request = bots_pb2.CreateBotSessionRequest(parent='',
-                                               bot_session=bot_session)
-    # Inject work
-    action_digest = remote_execution_pb2.Digest(hash='gaff')
-    _inject_work(instance._instances[""]._scheduler, action_digest=action_digest)
-
-    # Simulated the severed binding between client and server
-    response = copy.deepcopy(instance.CreateBotSession(request, context))
-
-    response.leases[0].state = LeaseState.ACTIVE.value
-
-    request = copy.deepcopy(bots_pb2.UpdateBotSessionRequest(name=response.name,
-                                                             bot_session=response))
-
-    response = instance.UpdateBotSession(request, context)
-
-    response.leases[0].state = state.value
-
-    request = bots_pb2.UpdateBotSessionRequest(name=response.name,
-                                               bot_session=response)
-
-    response = instance.UpdateBotSession(request, context)
-
-    context.set_code.assert_called_once_with(grpc.StatusCode.DATA_LOSS)
-
-
-def test_work_active_to_active(bot_session, context, instance):
-    request = bots_pb2.CreateBotSessionRequest(parent='',
-                                               bot_session=bot_session)
-    # Inject work
-    action_digest = remote_execution_pb2.Digest(hash='gaff')
-    _inject_work(instance._instances[""]._scheduler, action_digest=action_digest)
-
-    # Simulated the severed binding between client and server
-    response = copy.deepcopy(instance.CreateBotSession(request, context))
-
-    response.leases[0].state = LeaseState.ACTIVE.value
-
-    request = bots_pb2.UpdateBotSessionRequest(name=response.name,
-                                               bot_session=response)
-
-    response = instance.UpdateBotSession(request, context)
-
-    assert response.leases[0].state == LeaseState.ACTIVE.value
-
-
 def test_post_bot_event_temp(context, instance):
     request = bots_pb2.PostBotEventTempRequest()
     instance.PostBotEventTemp(request, context)

tests/integration/hardware_interface.py

+# Copyright (C) 2018 Bloomberg LP
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#  <http://www.apache.org/licenses/LICENSE-2.0>
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=redefined-outer-name
+
+
+import pytest
+
+from buildgrid._exceptions import FailedPreconditionError
+from buildgrid.bot.hardware.interface import HardwareInterface
+from buildgrid.bot.hardware.device import Device
+from buildgrid.bot.hardware.worker import Worker
+
+
+CONFIGURATIONS_WORKER = [{'DockerImage': ['Blade']}, {'DockerImage': ['Sheep', 'Snake']}]
+PROPERTIES_WORKER = [{'pool': ['Blade']}, {'pool': ['Sheep', 'Snake']}]
+PROPERTIES_DEVICE = [{'os': ['Blade']}, {'has-docker': ['Sheep', 'Snake']}]
+
+
+def test_create_hardware():
+    worker = Worker()
+    interface = HardwareInterface(worker)
+
+    device0 = Device()
+    worker.add_device(device0)
+    protobuf_worker = interface.get_worker_pb2()
+    assert len(protobuf_worker.devices) == 1
+
+    worker.add_device(Device())
+    protobuf_worker = interface.get_worker_pb2()
+    assert len(protobuf_worker.devices) == 2
+
+    assert protobuf_worker.devices[0].handle != protobuf_worker.devices[1].handle
+    assert device0.name == protobuf_worker.devices[0].handle
+
+
+@pytest.mark.parametrize('config', CONFIGURATIONS_WORKER)
+def test_worker_config(config):
+    worker = Worker(configs=config)
+    protobuf_worker = worker.get_pb2()
+
+    proto_cfg = {}
+    for cfg in protobuf_worker.configs:
+        k = cfg.key
+        v = cfg.value
+
+        proto_cfg_values = proto_cfg.get(k)
+        if not proto_cfg_values:
+            proto_cfg[k] = [v]
+        else:
+            proto_cfg_values.append(v)
+
+    assert config == proto_cfg
+    assert worker.configs == config
+
+
+@pytest.mark.parametrize('prop', PROPERTIES_WORKER)
+def test_worker_property(prop):
+    worker = Worker(properties=prop)
+    protobuf_worker = worker.get_pb2()
+
+    proto_prop = {}
+    for p in protobuf_worker.properties:
+        k = p.key
+        v = p.value
+
+        proto_prop_values = proto_prop.get(k)
+        if not proto_prop_values:
+            proto_prop[k] = [v]
+        else:
+            proto_prop_values.append(v)
+
+    assert prop == proto_prop
+    assert worker.properties == prop
+
+
+@pytest.mark.parametrize('prop', PROPERTIES_DEVICE)
+def test_device_property(prop):
+    device = Device(properties=prop)
+    protobuf_device = device.get_pb2()
+
+    proto_prop = {}
+    for p in protobuf_device.properties:
+        k = p.key
+        v = p.value
+
+        proto_prop_values = proto_prop.get(k)
+        if not proto_prop_values:
+            proto_prop[k] = [v]
+        else:
+            proto_prop_values.append(v)
+
+    assert prop == proto_prop
+    assert device.properties == prop
+
+
+@pytest.mark.parametrize('config', [{'piano': ['Blade']}])
+def test_worker_config_fail(config):
+    with pytest.raises(KeyError):
+        Worker(configs=config)
+
+
+@pytest.mark.parametrize('prop', [{'piano': ['Blade']}])
+def test_worker_property_fail(prop):
+    with pytest.raises(KeyError):
+        Worker(properties=prop)
+
+
+@pytest.mark.parametrize('prop', [{'piano': ['Blade']}])
+def test_device_property_fail(prop):
+    with pytest.raises(KeyError):
+        Device(properties=prop)
+
+
+@pytest.mark.parametrize('requirements', CONFIGURATIONS_WORKER)
+def test_configure_hardware(requirements):
+    hardware_interface = HardwareInterface(Worker(configs=requirements))
+
+    worker_requirements = Worker(configs=requirements)
+
+    hardware_interface.configure_hardware(worker_requirements.get_pb2())
+
+
+@pytest.mark.parametrize('requirements', CONFIGURATIONS_WORKER)
+def test_configure_hardware_fail(requirements):
+    hardware_interface = HardwareInterface(Worker())
+
+    worker_requirements = Worker(configs=requirements)
+
+    with pytest.raises(FailedPreconditionError):
+        hardware_interface.configure_hardware(worker_requirements.get_pb2())

tests/integration/operations_service.py

@@ -28,10 +28,8 @@ from buildgrid._enums import OperationStage
 from buildgrid._exceptions import InvalidArgumentError
 from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
 from buildgrid._protos.google.longrunning import operations_pb2
-from buildgrid._protos.google.rpc import status_pb2
 from buildgrid.server.cas.storage import lru_memory_cache
 from buildgrid.server.controller import ExecutionController
-from buildgrid.server.job import LeaseState
 from buildgrid.server.operations import service
 from buildgrid.server.operations.service import OperationsService
 from buildgrid.utils import create_digest
@@ -166,31 +164,6 @@ def test_list_operations_instance_fail(instance, controller, execute_request, co
     context.set_code.assert_called_once_with(grpc.StatusCode.INVALID_ARGUMENT)
 
 
-def test_list_operations_with_result(instance, controller, execute_request, context):
-    response_execute = controller.execution_instance.execute(execute_request.action_digest,
-                                                             execute_request.skip_cache_lookup)
-
-    action_result = remote_execution_pb2.ActionResult()
-    output_file = remote_execution_pb2.OutputFile(path='unicorn')
-    action_result.output_files.extend([output_file])
-
-    controller.operations_instance._scheduler.jobs[response_execute.name].create_lease()
-    controller.operations_instance._scheduler.update_job_lease_state(response_execute.name,
-                                                                     LeaseState.COMPLETED,
-                                                                     lease_status=status_pb2.Status(),
-                                                                     lease_result=_pack_any(action_result))
-
-    request = operations_pb2.ListOperationsRequest(name=instance_name)
-    response = instance.ListOperations(request, context)
-
-    assert response.operations[0].name.split('/')[-1] == response_execute.name
-
-    execute_response = remote_execution_pb2.ExecuteResponse()
-    response.operations[0].response.Unpack(execute_response)
-
-    assert execute_response.result.output_files == action_result.output_files
-
-
 def test_list_operations_empty(instance, context):
     request = operations_pb2.ListOperationsRequest(name=instance_name)
 

tests/utils/bots_interface.py

+# Copyright (C) 2018 Bloomberg LP
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#  <http://www.apache.org/licenses/LICENSE-2.0>
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from concurrent import futures
+from contextlib import contextmanager
+import multiprocessing
+import os
+import signal
+import uuid
+
+import grpc
+import pytest_cov
+
+from buildgrid._enums import LeaseState
+from buildgrid._protos.google.devtools.remoteworkers.v1test2 import bots_pb2
+from buildgrid.server.bots import service
+
+
+@contextmanager
+def serve_bots_interface(instances):
+    server = Server(instances)
+    try:
+        yield server
+    finally:
+        server.quit()
+
+
+# Use subprocess to avoid creation of gRPC threads in main process
+# See https://github.com/grpc/grpc/blob/master/doc/fork_support.md
+class Server:
+
+    def __init__(self, instances):
+        self.instances = instances
+
+        self.__queue = multiprocessing.Queue()
+        # Queue purely for bot session updates
+        self.__bot_session_queue = multiprocessing.Queue()
+        # Queue to send messages to subprocess
+        self.__message_queue = multiprocessing.Queue()
+        self.__process = multiprocessing.Process(
+            target=Server.serve,
+            args=(self.__queue, self.instances,
+                  self.__bot_session_queue, self.__message_queue))
+        self.__process.start()
+
+        self.port = self.__queue.get()
+        self.remote = 'localhost:{}'.format(self.port)
+
+    @staticmethod
+    def serve(queue, instances, bot_session_queue, message_queue):
+        pytest_cov.embed.cleanup_on_sigterm()
+
+        # Use max_workers default from Python 3.5+
+        max_workers = (os.cpu_count() or 1) * 5
+        server = grpc.server(futures.ThreadPoolExecutor(max_workers))
+        port = server.add_insecure_port('localhost:0')
+
+        bots_service = service.BotsService(server)
+        for name in instances:
+            bots_interface = BotsInterface(bot_session_queue, message_queue)
+            bots_service.add_instance(name, bots_interface)
+
+        server.start()
+        queue.put(port)
+        signal.pause()
+
+    def get_bot_session(self, timeout=1):
+        bot_session = bots_pb2.BotSession()
+        bot_session.ParseFromString(self.__bot_session_queue.get(timeout=timeout))
+        return bot_session
+
+    # Injects leases
+    def inject_work(self, lease=None, timeout=1):
+        if not lease:
+            lease = bots_pb2.Lease()
+            lease.state = LeaseState.PENDING.value
+
+        lease_string = lease.SerializeToString()
+        self.__message_queue.put(('INJECT_WORK', lease_string))
+
+    # Triggers a cancellation of a lease from server
+    def cancel_lease(self, lease_id):
+        self.__message_queue.put(('CANCEL_LEASE', lease_id))
+
+    def quit(self):
+        if self.__process:
+            self.__process.terminate()
+            self.__process.join()
+
+
+class BotsInterface:
+
+    def __init__(self, bot_session_queue, message_queue):
+        self.__bot_session_queue = bot_session_queue
+        self.__message_queue = message_queue
+
+    def register_instance_with_server(self, instance_name, server):
+        server.add_bots_interface(self, instance_name)
+
+    def create_bot_session(self, parent, bot_session):
+        name = "{}/{}".format(parent, str(uuid.uuid4()))
+        bot_session.name = name
+
+        while not self.__message_queue.empty():
+            message = self.__message_queue.get()
+            if message[0] == 'INJECT_WORK':
+                lease_string = message[1]
+                lease = bots_pb2.Lease()
+                lease.ParseFromString(lease_string)
+                bot_session.leases.extend([lease])
+
+        self.__bot_session_queue.put(bot_session.SerializeToString())
+        return bot_session
+
+    def update_bot_session(self, name, bot_session):
+        for lease in bot_session.leases:
+            state = LeaseState(lease.state)
+            if state == LeaseState.COMPLETED:
+                lease.Clear()
+
+            elif state == LeaseState.CANCELLED:
+                lease.Clear()
+
+        while not self.__message_queue.empty():
+            message = self.__message_queue.get()
+
+            if message[0] == 'INJECT_WORK':
+                lease_string = message[1]
+                lease = bots_pb2.Lease()
+                lease.ParseFromString(lease_string)
+                bot_session.leases.extend([lease])
+
+            elif message[0] == 'CANCEL_LEASE':
+                lease_id = message[1]
+                for lease in bot_session.leases:
+                    if lease.id == lease_id:
+                        lease.state = LeaseState.CANCELLED.value
+
+        self.__bot_session_queue.put(bot_session.SerializeToString())
+        return bot_session