finn pushed to branch finn/48-cancellation-leases at BuildGrid / buildgrid
Commits:
- 
22fef311
by Finn at 2018-11-23T11:59:17Z
- 
af723c5f
by Finn at 2018-11-23T11:59:17Z
- 
40c0db97
by Finn at 2018-11-23T11:59:17Z
- 
61629086
by Finn at 2018-11-23T11:59:17Z
- 
28611d0f
by Finn at 2018-11-23T11:59:17Z
- 
7899d14a
by Finn at 2018-11-23T11:59:17Z
7 changed files:
- tests/cas/test_client.py
- tests/cas/test_storage.py
- tests/integration/bot_session.py
- tests/integration/bots_service.py
- + tests/integration/hardware_interface.py
- tests/integration/operations_service.py
- + tests/utils/bots_interface.py
Changes:
| ... | ... | @@ -26,7 +26,8 @@ from buildgrid.client.cas import download, upload | 
| 26 | 26 |  from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
 | 
| 27 | 27 |  from buildgrid.utils import create_digest
 | 
| 28 | 28 |  | 
| 29 | -from ..utils.cas import serve_cas, run_in_subprocess
 | |
| 29 | +from ..utils.cas import serve_cas
 | |
| 30 | +from ..utils.utils import run_in_subprocess
 | |
| 30 | 31 |  | 
| 31 | 32 |  | 
| 32 | 33 |  INTANCES = ['', 'instance']
 | 
| ... | ... | @@ -32,7 +32,8 @@ from buildgrid.server.cas.storage.s3 import S3Storage | 
| 32 | 32 |  from buildgrid.server.cas.storage.with_cache import WithCacheStorage
 | 
| 33 | 33 |  from buildgrid.settings import HASH
 | 
| 34 | 34 |  | 
| 35 | -from ..utils.cas import serve_cas, run_in_subprocess
 | |
| 35 | +from ..utils.cas import serve_cas
 | |
| 36 | +from ..utils.utils import run_in_subprocess
 | |
| 36 | 37 |  | 
| 37 | 38 |  | 
| 38 | 39 |  BLOBS = [(b'abc', b'defg', b'hijk', b'')]
 | 
| ... | ... | @@ -14,56 +14,174 @@ | 
| 14 | 14 |  | 
| 15 | 15 |  # pylint: disable=redefined-outer-name
 | 
| 16 | 16 |  | 
| 17 | -import uuid
 | |
| 17 | +import asyncio
 | |
| 18 | 18 |  | 
| 19 | +import grpc
 | |
| 19 | 20 |  import pytest
 | 
| 20 | 21 |  | 
| 21 | -from buildgrid.bot import bot_session
 | |
| 22 | +from buildgrid._enums import LeaseState
 | |
| 23 | +from buildgrid._protos.google.devtools.remoteworkers.v1test2 import bots_pb2
 | |
| 24 | +from buildgrid.bot.hardware.worker import Worker
 | |
| 25 | +from buildgrid.bot.hardware.interface import HardwareInterface
 | |
| 26 | +from buildgrid.bot.session import BotSession
 | |
| 27 | +from buildgrid.bot.interface import BotInterface
 | |
| 22 | 28 |  | 
| 29 | +from ..utils.utils import run_in_subprocess
 | |
| 30 | +from ..utils.bots_interface import serve_bots_interface
 | |
| 23 | 31 |  | 
| 24 | -@pytest.mark.parametrize("docker_value", ["True", "False"])
 | |
| 25 | -@pytest.mark.parametrize("os_value", ["nexus7", "nexus8"])
 | |
| 26 | -def test_create_device(docker_value, os_value):
 | |
| 27 | -    properties = {'docker': docker_value, 'os': os_value}
 | |
| 28 | -    device = bot_session.Device(properties)
 | |
| 29 | 32 |  | 
| 30 | -    assert uuid.UUID(device.name, version=4)
 | |
| 31 | -    assert properties == device.properties
 | |
| 33 | +INSTANCES = ['', 'instance']
 | |
| 32 | 34 |  | 
| 33 | 35 |  | 
| 34 | -def test_create_device_key_fail():
 | |
| 35 | -    properties = {'voight': 'kampff'}
 | |
| 36 | +# Use subprocess to avoid creation of gRPC threads in main process
 | |
| 37 | +# See https://github.com/grpc/grpc/blob/master/doc/fork_support.md
 | |
| 38 | +# Multiprocessing uses pickle which protobufs don't work with
 | |
| 39 | +# Workaround wrapper to send messages as strings
 | |
| 40 | +class ServerInterface:
 | |
| 36 | 41 |  | 
| 37 | -    with pytest.raises(KeyError):
 | |
| 38 | -        bot_session.Device(properties)
 | |
| 42 | +    def __init__(self, remote):
 | |
| 43 | +        self.__remote = remote
 | |
| 39 | 44 |  | 
| 45 | +    def create_bot_session(self, parent, bot_session):
 | |
| 40 | 46 |  | 
| 41 | -def test_create_device_value_fail():
 | |
| 42 | -    properties = {'docker': True}
 | |
| 47 | +        def __create_bot_session(queue, remote, parent, string_bot_session):
 | |
| 48 | +            bot_session = bots_pb2.BotSession()
 | |
| 49 | +            bot_session.ParseFromString(string_bot_session)
 | |
| 43 | 50 |  | 
| 44 | -    with pytest.raises(ValueError):
 | |
| 45 | -        bot_session.Device(properties)
 | |
| 51 | +            interface = BotInterface(grpc.insecure_channel(remote))
 | |
| 46 | 52 |  | 
| 53 | +            result = interface.create_bot_session(parent, bot_session)
 | |
| 54 | +            queue.put(result.SerializeToString())
 | |
| 47 | 55 |  | 
| 48 | -def test_create_worker():
 | |
| 49 | -    properties = {'pool': 'swim'}
 | |
| 50 | -    configs = {'DockerImage': 'Windows'}
 | |
| 51 | -    worker = bot_session.Worker(properties, configs)
 | |
| 56 | +        string_bot_session = bot_session.SerializeToString()
 | |
| 57 | +        result = run_in_subprocess(__create_bot_session,
 | |
| 58 | +                                   self.__remote, parent, string_bot_session)
 | |
| 52 | 59 |  | 
| 53 | -    assert properties == worker.properties
 | |
| 54 | -    assert configs == worker.configs
 | |
| 60 | +        bot_session = bots_pb2.BotSession()
 | |
| 61 | +        bot_session.ParseFromString(result)
 | |
| 62 | +        return bot_session
 | |
| 55 | 63 |  | 
| 56 | -    device = bot_session.Device()
 | |
| 57 | -    worker.add_device(device)
 | |
| 64 | +    def update_bot_session(self, bot_session, update_mask=None):
 | |
| 58 | 65 |  | 
| 59 | -    assert worker._devices[0] == device
 | |
| 66 | +        def __update_bot_session(queue, remote, string_bot_session, update_mask):
 | |
| 67 | +            bot_session = bots_pb2.BotSession()
 | |
| 68 | +            bot_session.ParseFromString(string_bot_session)
 | |
| 60 | 69 |  | 
| 70 | +            interface = BotInterface(grpc.insecure_channel(remote))
 | |
| 61 | 71 |  | 
| 62 | -def test_create_worker_key_fail():
 | |
| 63 | -    properties = {'voight': 'kampff'}
 | |
| 64 | -    configs = {'voight': 'kampff'}
 | |
| 72 | +            result = interface.update_bot_session(bot_session, update_mask)
 | |
| 73 | +            queue.put(result.SerializeToString())
 | |
| 65 | 74 |  | 
| 66 | -    with pytest.raises(KeyError):
 | |
| 67 | -        bot_session.Worker(properties)
 | |
| 68 | -    with pytest.raises(KeyError):
 | |
| 69 | -        bot_session.Worker(configs) | |
| 75 | +        string_bot_session = bot_session.SerializeToString()
 | |
| 76 | +        result = run_in_subprocess(__update_bot_session,
 | |
| 77 | +                                   self.__remote, string_bot_session, update_mask)
 | |
| 78 | + | |
| 79 | +        bot_session = bots_pb2.BotSession()
 | |
| 80 | +        bot_session.ParseFromString(result)
 | |
| 81 | +        return bot_session
 | |
| 82 | + | |
| 83 | + | |
| 84 | +@pytest.mark.parametrize('instance', INSTANCES)
 | |
| 85 | +def test_create_bot_session(instance):
 | |
| 86 | + | |
| 87 | +    with serve_bots_interface([instance]) as server:
 | |
| 88 | +        interface = ServerInterface(server.remote)
 | |
| 89 | +        hardware_interface = HardwareInterface(Worker())
 | |
| 90 | +        session = BotSession(instance, interface, hardware_interface, None)
 | |
| 91 | +        session.create_bot_session()
 | |
| 92 | +        assert session.get_pb2() == server.get_bot_session()
 | |
| 93 | + | |
| 94 | + | |
| 95 | +@pytest.mark.parametrize('instance', INSTANCES)
 | |
| 96 | +def test_update_bot_session(instance):
 | |
| 97 | + | |
| 98 | +    with serve_bots_interface([instance]) as server:
 | |
| 99 | +        interface = ServerInterface(server.remote)
 | |
| 100 | +        hardware_interface = HardwareInterface(Worker())
 | |
| 101 | +        session = BotSession(instance, interface, hardware_interface, None)
 | |
| 102 | +        session.create_bot_session()
 | |
| 103 | +        assert session.get_pb2() == server.get_bot_session()
 | |
| 104 | +        session.update_bot_session()
 | |
| 105 | +        assert session.get_pb2() == server.get_bot_session()
 | |
| 106 | + | |
| 107 | + | |
| 108 | +@pytest.mark.parametrize('instance', INSTANCES)
 | |
| 109 | +def test_create_bot_session_with_work(instance):
 | |
| 110 | + | |
| 111 | +    def __work(lease, context, event):
 | |
| 112 | +        return lease
 | |
| 113 | + | |
| 114 | +    with serve_bots_interface([instance]) as server:
 | |
| 115 | +        interface = ServerInterface(server.remote)
 | |
| 116 | +        hardware_interface = HardwareInterface(Worker())
 | |
| 117 | +        session = BotSession(instance, interface, hardware_interface, __work)
 | |
| 118 | +        server.inject_work()
 | |
| 119 | +        session.create_bot_session()
 | |
| 120 | + | |
| 121 | +        assert len(session.get_pb2().leases) == 1
 | |
| 122 | + | |
| 123 | +        loop = asyncio.get_event_loop()
 | |
| 124 | +        for task in asyncio.Task.all_tasks():
 | |
| 125 | +            loop.run_until_complete(task)
 | |
| 126 | + | |
| 127 | +        assert session.get_pb2().leases[0].state == LeaseState.COMPLETED.value
 | |
| 128 | + | |
| 129 | + | |
| 130 | +@pytest.mark.parametrize('instance', INSTANCES)
 | |
| 131 | +def test_update_bot_session_with_work(instance):
 | |
| 132 | + | |
| 133 | +    def __work(lease, context, event):
 | |
| 134 | +        return lease
 | |
| 135 | + | |
| 136 | +    with serve_bots_interface([instance]) as server:
 | |
| 137 | +        interface = ServerInterface(server.remote)
 | |
| 138 | +        hardware_interface = HardwareInterface(Worker())
 | |
| 139 | +        session = BotSession(instance, interface, hardware_interface, __work)
 | |
| 140 | +        session.create_bot_session()
 | |
| 141 | +        server.inject_work()
 | |
| 142 | +        session.update_bot_session()
 | |
| 143 | + | |
| 144 | +        assert len(session.get_pb2().leases) == 1
 | |
| 145 | + | |
| 146 | +        loop = asyncio.get_event_loop()
 | |
| 147 | +        for task in asyncio.Task.all_tasks():
 | |
| 148 | +            loop.run_until_complete(task)
 | |
| 149 | + | |
| 150 | +        assert session.get_pb2().leases[0].state == LeaseState.COMPLETED.value
 | |
| 151 | + | |
| 152 | + | |
| 153 | +@pytest.mark.parametrize('instance', INSTANCES)
 | |
| 154 | +def test_cancel_leases(instance):
 | |
| 155 | + | |
| 156 | +    def __work(lease, context, cancel_event):
 | |
| 157 | +        # while not cancel_event.is_set():
 | |
| 158 | + | |
| 159 | +        return lease
 | |
| 160 | + | |
| 161 | +    with serve_bots_interface([instance]) as server:
 | |
| 162 | +        interface = ServerInterface(server.remote)
 | |
| 163 | +        hardware_interface = HardwareInterface(Worker())
 | |
| 164 | +        session = BotSession(instance, interface, hardware_interface, __work)
 | |
| 165 | + | |
| 166 | +        lease = bots_pb2.Lease()
 | |
| 167 | +        lease.state = LeaseState.PENDING.value
 | |
| 168 | +        lease.id = 'foo'
 | |
| 169 | +        server.inject_work(lease)
 | |
| 170 | +        session.create_bot_session()
 | |
| 171 | + | |
| 172 | +        leases_pb2 = session.get_pb2().leases
 | |
| 173 | +        assert len(leases_pb2) == 1
 | |
| 174 | +        assert leases_pb2[0].state == LeaseState.ACTIVE.value
 | |
| 175 | + | |
| 176 | +        server.cancel_lease(leases_pb2[0].id)
 | |
| 177 | +        session.update_bot_session()
 | |
| 178 | +        assert len(session.get_pb2().leases) == 1
 | |
| 179 | + | |
| 180 | +        loop = asyncio.get_event_loop()
 | |
| 181 | +        for task in asyncio.Task.all_tasks():
 | |
| 182 | +            try:
 | |
| 183 | +                loop.run_until_complete(task)
 | |
| 184 | +            except asyncio.CancelledError:
 | |
| 185 | +                pass
 | |
| 186 | + | |
| 187 | +        assert session.get_pb2().leases[0].state == LeaseState.CANCELLED.value | 
| ... | ... | @@ -17,7 +17,6 @@ | 
| 17 | 17 |  | 
| 18 | 18 |  # pylint: disable=redefined-outer-name
 | 
| 19 | 19 |  | 
| 20 | -import copy
 | |
| 21 | 20 |  from unittest import mock
 | 
| 22 | 21 |  | 
| 23 | 22 |  import grpc
 | 
| ... | ... | @@ -150,129 +149,6 @@ def test_update_leases_with_work(bot_session, context, instance): | 
| 150 | 149 |      assert response_action == action_digest
 | 
| 151 | 150 |  | 
| 152 | 151 |  | 
| 153 | -def test_update_leases_work_complete(bot_session, context, instance):
 | |
| 154 | -    request = bots_pb2.CreateBotSessionRequest(parent='',
 | |
| 155 | -                                               bot_session=bot_session)
 | |
| 156 | -    # Create bot session
 | |
| 157 | -    # Simulated the severed binding between client and server
 | |
| 158 | -    response = copy.deepcopy(instance.CreateBotSession(request, context))
 | |
| 159 | - | |
| 160 | -    # Inject work
 | |
| 161 | -    action_digest = remote_execution_pb2.Digest(hash='gaff')
 | |
| 162 | -    _inject_work(instance._instances[""]._scheduler, action_digest=action_digest)
 | |
| 163 | - | |
| 164 | -    request = bots_pb2.UpdateBotSessionRequest(name=response.name,
 | |
| 165 | -                                               bot_session=response)
 | |
| 166 | -    response = copy.deepcopy(instance.UpdateBotSession(request, context))
 | |
| 167 | - | |
| 168 | -    assert response.leases[0].state == LeaseState.PENDING.value
 | |
| 169 | -    response.leases[0].state = LeaseState.ACTIVE.value
 | |
| 170 | - | |
| 171 | -    request = bots_pb2.UpdateBotSessionRequest(name=response.name,
 | |
| 172 | -                                               bot_session=response)
 | |
| 173 | - | |
| 174 | -    response = copy.deepcopy(instance.UpdateBotSession(request, context))
 | |
| 175 | - | |
| 176 | -    response.leases[0].state = LeaseState.COMPLETED.value
 | |
| 177 | -    response.leases[0].result.Pack(remote_execution_pb2.ActionResult())
 | |
| 178 | - | |
| 179 | -    request = bots_pb2.UpdateBotSessionRequest(name=response.name,
 | |
| 180 | -                                               bot_session=response)
 | |
| 181 | -    response = copy.deepcopy(instance.UpdateBotSession(request, context))
 | |
| 182 | - | |
| 183 | -    assert len(response.leases) is 0
 | |
| 184 | - | |
| 185 | - | |
| 186 | -def test_work_rejected_by_bot(bot_session, context, instance):
 | |
| 187 | -    request = bots_pb2.CreateBotSessionRequest(parent='',
 | |
| 188 | -                                               bot_session=bot_session)
 | |
| 189 | -    # Inject work
 | |
| 190 | -    action_digest = remote_execution_pb2.Digest(hash='gaff')
 | |
| 191 | -    _inject_work(instance._instances[""]._scheduler, action_digest=action_digest)
 | |
| 192 | - | |
| 193 | -    # Simulated the severed binding between client and server
 | |
| 194 | -    response = copy.deepcopy(instance.CreateBotSession(request, context))
 | |
| 195 | - | |
| 196 | -    # Reject work
 | |
| 197 | -    assert response.leases[0].state == LeaseState.PENDING.value
 | |
| 198 | -    response.leases[0].state = LeaseState.COMPLETED.value
 | |
| 199 | -    request = bots_pb2.UpdateBotSessionRequest(name=response.name,
 | |
| 200 | -                                               bot_session=response)
 | |
| 201 | - | |
| 202 | -    response = instance.UpdateBotSession(request, context)
 | |
| 203 | - | |
| 204 | -    context.set_code.assert_called_once_with(grpc.StatusCode.UNIMPLEMENTED)
 | |
| 205 | - | |
| 206 | - | |
| 207 | -@pytest.mark.parametrize("state", [LeaseState.LEASE_STATE_UNSPECIFIED, LeaseState.PENDING])
 | |
| 208 | -def test_work_out_of_sync_from_pending(state, bot_session, context, instance):
 | |
| 209 | -    request = bots_pb2.CreateBotSessionRequest(parent='',
 | |
| 210 | -                                               bot_session=bot_session)
 | |
| 211 | -    # Inject work
 | |
| 212 | -    action_digest = remote_execution_pb2.Digest(hash='gaff')
 | |
| 213 | -    _inject_work(instance._instances[""]._scheduler, action_digest=action_digest)
 | |
| 214 | - | |
| 215 | -    # Simulated the severed binding between client and server
 | |
| 216 | -    response = copy.deepcopy(instance.CreateBotSession(request, context))
 | |
| 217 | - | |
| 218 | -    response.leases[0].state = state.value
 | |
| 219 | - | |
| 220 | -    request = bots_pb2.UpdateBotSessionRequest(name=response.name,
 | |
| 221 | -                                               bot_session=response)
 | |
| 222 | - | |
| 223 | -    response = instance.UpdateBotSession(request, context)
 | |
| 224 | - | |
| 225 | -    context.set_code.assert_called_once_with(grpc.StatusCode.DATA_LOSS)
 | |
| 226 | - | |
| 227 | - | |
| 228 | -@pytest.mark.parametrize("state", [LeaseState.LEASE_STATE_UNSPECIFIED, LeaseState.PENDING])
 | |
| 229 | -def test_work_out_of_sync_from_active(state, bot_session, context, instance):
 | |
| 230 | -    request = bots_pb2.CreateBotSessionRequest(parent='',
 | |
| 231 | -                                               bot_session=bot_session)
 | |
| 232 | -    # Inject work
 | |
| 233 | -    action_digest = remote_execution_pb2.Digest(hash='gaff')
 | |
| 234 | -    _inject_work(instance._instances[""]._scheduler, action_digest=action_digest)
 | |
| 235 | - | |
| 236 | -    # Simulated the severed binding between client and server
 | |
| 237 | -    response = copy.deepcopy(instance.CreateBotSession(request, context))
 | |
| 238 | - | |
| 239 | -    response.leases[0].state = LeaseState.ACTIVE.value
 | |
| 240 | - | |
| 241 | -    request = copy.deepcopy(bots_pb2.UpdateBotSessionRequest(name=response.name,
 | |
| 242 | -                                                             bot_session=response))
 | |
| 243 | - | |
| 244 | -    response = instance.UpdateBotSession(request, context)
 | |
| 245 | - | |
| 246 | -    response.leases[0].state = state.value
 | |
| 247 | - | |
| 248 | -    request = bots_pb2.UpdateBotSessionRequest(name=response.name,
 | |
| 249 | -                                               bot_session=response)
 | |
| 250 | - | |
| 251 | -    response = instance.UpdateBotSession(request, context)
 | |
| 252 | - | |
| 253 | -    context.set_code.assert_called_once_with(grpc.StatusCode.DATA_LOSS)
 | |
| 254 | - | |
| 255 | - | |
| 256 | -def test_work_active_to_active(bot_session, context, instance):
 | |
| 257 | -    request = bots_pb2.CreateBotSessionRequest(parent='',
 | |
| 258 | -                                               bot_session=bot_session)
 | |
| 259 | -    # Inject work
 | |
| 260 | -    action_digest = remote_execution_pb2.Digest(hash='gaff')
 | |
| 261 | -    _inject_work(instance._instances[""]._scheduler, action_digest=action_digest)
 | |
| 262 | - | |
| 263 | -    # Simulated the severed binding between client and server
 | |
| 264 | -    response = copy.deepcopy(instance.CreateBotSession(request, context))
 | |
| 265 | - | |
| 266 | -    response.leases[0].state = LeaseState.ACTIVE.value
 | |
| 267 | - | |
| 268 | -    request = bots_pb2.UpdateBotSessionRequest(name=response.name,
 | |
| 269 | -                                               bot_session=response)
 | |
| 270 | - | |
| 271 | -    response = instance.UpdateBotSession(request, context)
 | |
| 272 | - | |
| 273 | -    assert response.leases[0].state == LeaseState.ACTIVE.value
 | |
| 274 | - | |
| 275 | - | |
| 276 | 152 |  def test_post_bot_event_temp(context, instance):
 | 
| 277 | 153 |      request = bots_pb2.PostBotEventTempRequest()
 | 
| 278 | 154 |      instance.PostBotEventTemp(request, context)
 | 
| 1 | +# Copyright (C) 2018 Bloomberg LP
 | |
| 2 | +#
 | |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License");
 | |
| 4 | +# you may not use this file except in compliance with the License.
 | |
| 5 | +# You may obtain a copy of the License at
 | |
| 6 | +#
 | |
| 7 | +#  <http://www.apache.org/licenses/LICENSE-2.0>
 | |
| 8 | +#
 | |
| 9 | +# Unless required by applicable law or agreed to in writing, software
 | |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS,
 | |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 | |
| 12 | +# See the License for the specific language governing permissions and
 | |
| 13 | +# limitations under the License.
 | |
| 14 | + | |
| 15 | +# pylint: disable=redefined-outer-name
 | |
| 16 | + | |
| 17 | + | |
| 18 | +import pytest
 | |
| 19 | + | |
| 20 | +from buildgrid._exceptions import FailedPreconditionError
 | |
| 21 | +from buildgrid.bot.hardware.interface import HardwareInterface
 | |
| 22 | +from buildgrid.bot.hardware.device import Device
 | |
| 23 | +from buildgrid.bot.hardware.worker import Worker
 | |
| 24 | + | |
| 25 | + | |
| 26 | +CONFIGURATIONS_WORKER = [{'DockerImage': ['Blade']}, {'DockerImage': ['Sheep', 'Snake']}]
 | |
| 27 | +PROPERTIES_WORKER = [{'pool': ['Blade']}, {'pool': ['Sheep', 'Snake']}]
 | |
| 28 | +PROPERTIES_DEVICE = [{'os': ['Blade']}, {'has-docker': ['Sheep', 'Snake']}]
 | |
| 29 | + | |
| 30 | + | |
| 31 | +def test_create_hardware():
 | |
| 32 | +    worker = Worker()
 | |
| 33 | +    interface = HardwareInterface(worker)
 | |
| 34 | + | |
| 35 | +    device0 = Device()
 | |
| 36 | +    worker.add_device(device0)
 | |
| 37 | +    protobuf_worker = interface.get_worker_pb2()
 | |
| 38 | +    assert len(protobuf_worker.devices) == 1
 | |
| 39 | + | |
| 40 | +    worker.add_device(Device())
 | |
| 41 | +    protobuf_worker = interface.get_worker_pb2()
 | |
| 42 | +    assert len(protobuf_worker.devices) == 2
 | |
| 43 | + | |
| 44 | +    assert protobuf_worker.devices[0].handle != protobuf_worker.devices[1].handle
 | |
| 45 | +    assert device0.name == protobuf_worker.devices[0].handle
 | |
| 46 | + | |
| 47 | + | |
| 48 | +@pytest.mark.parametrize('config', CONFIGURATIONS_WORKER)
 | |
| 49 | +def test_worker_config(config):
 | |
| 50 | +    worker = Worker(configs=config)
 | |
| 51 | +    protobuf_worker = worker.get_pb2()
 | |
| 52 | + | |
| 53 | +    proto_cfg = {}
 | |
| 54 | +    for cfg in protobuf_worker.configs:
 | |
| 55 | +        k = cfg.key
 | |
| 56 | +        v = cfg.value
 | |
| 57 | + | |
| 58 | +        proto_cfg_values = proto_cfg.get(k)
 | |
| 59 | +        if not proto_cfg_values:
 | |
| 60 | +            proto_cfg[k] = [v]
 | |
| 61 | +        else:
 | |
| 62 | +            proto_cfg_values.append(v)
 | |
| 63 | + | |
| 64 | +    assert config == proto_cfg
 | |
| 65 | +    assert worker.configs == config
 | |
| 66 | + | |
| 67 | + | |
| 68 | +@pytest.mark.parametrize('prop', PROPERTIES_WORKER)
 | |
| 69 | +def test_worker_property(prop):
 | |
| 70 | +    worker = Worker(properties=prop)
 | |
| 71 | +    protobuf_worker = worker.get_pb2()
 | |
| 72 | + | |
| 73 | +    proto_prop = {}
 | |
| 74 | +    for p in protobuf_worker.properties:
 | |
| 75 | +        k = p.key
 | |
| 76 | +        v = p.value
 | |
| 77 | + | |
| 78 | +        proto_prop_values = proto_prop.get(k)
 | |
| 79 | +        if not proto_prop_values:
 | |
| 80 | +            proto_prop[k] = [v]
 | |
| 81 | +        else:
 | |
| 82 | +            proto_prop_values.append(v)
 | |
| 83 | + | |
| 84 | +    assert prop == proto_prop
 | |
| 85 | +    assert worker.properties == prop
 | |
| 86 | + | |
| 87 | + | |
| 88 | +@pytest.mark.parametrize('prop', PROPERTIES_DEVICE)
 | |
| 89 | +def test_device_property(prop):
 | |
| 90 | +    device = Device(properties=prop)
 | |
| 91 | +    protobuf_device = device.get_pb2()
 | |
| 92 | + | |
| 93 | +    proto_prop = {}
 | |
| 94 | +    for p in protobuf_device.properties:
 | |
| 95 | +        k = p.key
 | |
| 96 | +        v = p.value
 | |
| 97 | + | |
| 98 | +        proto_prop_values = proto_prop.get(k)
 | |
| 99 | +        if not proto_prop_values:
 | |
| 100 | +            proto_prop[k] = [v]
 | |
| 101 | +        else:
 | |
| 102 | +            proto_prop_values.append(v)
 | |
| 103 | + | |
| 104 | +    assert prop == proto_prop
 | |
| 105 | +    assert device.properties == prop
 | |
| 106 | + | |
| 107 | + | |
| 108 | +@pytest.mark.parametrize('config', [{'piano': ['Blade']}])
 | |
| 109 | +def test_worker_config_fail(config):
 | |
| 110 | +    with pytest.raises(KeyError):
 | |
| 111 | +        Worker(configs=config)
 | |
| 112 | + | |
| 113 | + | |
| 114 | +@pytest.mark.parametrize('prop', [{'piano': ['Blade']}])
 | |
| 115 | +def test_worker_property_fail(prop):
 | |
| 116 | +    with pytest.raises(KeyError):
 | |
| 117 | +        Worker(properties=prop)
 | |
| 118 | + | |
| 119 | + | |
| 120 | +@pytest.mark.parametrize('prop', [{'piano': ['Blade']}])
 | |
| 121 | +def test_device_property_fail(prop):
 | |
| 122 | +    with pytest.raises(KeyError):
 | |
| 123 | +        Device(properties=prop)
 | |
| 124 | + | |
| 125 | + | |
| 126 | +@pytest.mark.parametrize('requirements', CONFIGURATIONS_WORKER)
 | |
| 127 | +def test_configure_hardware(requirements):
 | |
| 128 | +    hardware_interface = HardwareInterface(Worker(configs=requirements))
 | |
| 129 | + | |
| 130 | +    worker_requirements = Worker(configs=requirements)
 | |
| 131 | + | |
| 132 | +    hardware_interface.configure_hardware(worker_requirements.get_pb2())
 | |
| 133 | + | |
| 134 | + | |
| 135 | +@pytest.mark.parametrize('requirements', CONFIGURATIONS_WORKER)
 | |
| 136 | +def test_configure_hardware_fail(requirements):
 | |
| 137 | +    hardware_interface = HardwareInterface(Worker())
 | |
| 138 | + | |
| 139 | +    worker_requirements = Worker(configs=requirements)
 | |
| 140 | + | |
| 141 | +    with pytest.raises(FailedPreconditionError):
 | |
| 142 | +        hardware_interface.configure_hardware(worker_requirements.get_pb2()) | 
| ... | ... | @@ -28,10 +28,8 @@ from buildgrid._enums import OperationStage | 
| 28 | 28 |  from buildgrid._exceptions import InvalidArgumentError
 | 
| 29 | 29 |  from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
 | 
| 30 | 30 |  from buildgrid._protos.google.longrunning import operations_pb2
 | 
| 31 | -from buildgrid._protos.google.rpc import status_pb2
 | |
| 32 | 31 |  from buildgrid.server.cas.storage import lru_memory_cache
 | 
| 33 | 32 |  from buildgrid.server.controller import ExecutionController
 | 
| 34 | -from buildgrid.server.job import LeaseState
 | |
| 35 | 33 |  from buildgrid.server.operations import service
 | 
| 36 | 34 |  from buildgrid.server.operations.service import OperationsService
 | 
| 37 | 35 |  from buildgrid.utils import create_digest
 | 
| ... | ... | @@ -166,31 +164,6 @@ def test_list_operations_instance_fail(instance, controller, execute_request, co | 
| 166 | 164 |      context.set_code.assert_called_once_with(grpc.StatusCode.INVALID_ARGUMENT)
 | 
| 167 | 165 |  | 
| 168 | 166 |  | 
| 169 | -def test_list_operations_with_result(instance, controller, execute_request, context):
 | |
| 170 | -    response_execute = controller.execution_instance.execute(execute_request.action_digest,
 | |
| 171 | -                                                             execute_request.skip_cache_lookup)
 | |
| 172 | - | |
| 173 | -    action_result = remote_execution_pb2.ActionResult()
 | |
| 174 | -    output_file = remote_execution_pb2.OutputFile(path='unicorn')
 | |
| 175 | -    action_result.output_files.extend([output_file])
 | |
| 176 | - | |
| 177 | -    controller.operations_instance._scheduler.jobs[response_execute.name].create_lease()
 | |
| 178 | -    controller.operations_instance._scheduler.update_job_lease_state(response_execute.name,
 | |
| 179 | -                                                                     LeaseState.COMPLETED,
 | |
| 180 | -                                                                     lease_status=status_pb2.Status(),
 | |
| 181 | -                                                                     lease_result=_pack_any(action_result))
 | |
| 182 | - | |
| 183 | -    request = operations_pb2.ListOperationsRequest(name=instance_name)
 | |
| 184 | -    response = instance.ListOperations(request, context)
 | |
| 185 | - | |
| 186 | -    assert response.operations[0].name.split('/')[-1] == response_execute.name
 | |
| 187 | - | |
| 188 | -    execute_response = remote_execution_pb2.ExecuteResponse()
 | |
| 189 | -    response.operations[0].response.Unpack(execute_response)
 | |
| 190 | - | |
| 191 | -    assert execute_response.result.output_files == action_result.output_files
 | |
| 192 | - | |
| 193 | - | |
| 194 | 167 |  def test_list_operations_empty(instance, context):
 | 
| 195 | 168 |      request = operations_pb2.ListOperationsRequest(name=instance_name)
 | 
| 196 | 169 |  | 
| 1 | +# Copyright (C) 2018 Bloomberg LP
 | |
| 2 | +#
 | |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License");
 | |
| 4 | +# you may not use this file except in compliance with the License.
 | |
| 5 | +# You may obtain a copy of the License at
 | |
| 6 | +#
 | |
| 7 | +#  <http://www.apache.org/licenses/LICENSE-2.0>
 | |
| 8 | +#
 | |
| 9 | +# Unless required by applicable law or agreed to in writing, software
 | |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS,
 | |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 | |
| 12 | +# See the License for the specific language governing permissions and
 | |
| 13 | +# limitations under the License.
 | |
| 14 | + | |
| 15 | +from concurrent import futures
 | |
| 16 | +from contextlib import contextmanager
 | |
| 17 | +import multiprocessing
 | |
| 18 | +import os
 | |
| 19 | +import signal
 | |
| 20 | +import uuid
 | |
| 21 | + | |
| 22 | +import grpc
 | |
| 23 | +import pytest_cov
 | |
| 24 | + | |
| 25 | +from buildgrid._enums import LeaseState
 | |
| 26 | +from buildgrid._protos.google.devtools.remoteworkers.v1test2 import bots_pb2
 | |
| 27 | +from buildgrid.server.bots import service
 | |
| 28 | + | |
| 29 | + | |
| 30 | +@contextmanager
 | |
| 31 | +def serve_bots_interface(instances):
 | |
| 32 | +    server = Server(instances)
 | |
| 33 | +    try:
 | |
| 34 | +        yield server
 | |
| 35 | +    finally:
 | |
| 36 | +        server.quit()
 | |
| 37 | + | |
| 38 | + | |
| 39 | +# Use subprocess to avoid creation of gRPC threads in main process
 | |
| 40 | +# See https://github.com/grpc/grpc/blob/master/doc/fork_support.md
 | |
| 41 | +class Server:
 | |
| 42 | + | |
| 43 | +    def __init__(self, instances):
 | |
| 44 | +        self.instances = instances
 | |
| 45 | + | |
| 46 | +        self.__queue = multiprocessing.Queue()
 | |
| 47 | +        # Queue purely for bot session updates
 | |
| 48 | +        self.__bot_session_queue = multiprocessing.Queue()
 | |
| 49 | +        # Queue to send messages to subprocess
 | |
| 50 | +        self.__message_queue = multiprocessing.Queue()
 | |
| 51 | +        self.__process = multiprocessing.Process(
 | |
| 52 | +            target=Server.serve,
 | |
| 53 | +            args=(self.__queue, self.instances,
 | |
| 54 | +                  self.__bot_session_queue, self.__message_queue))
 | |
| 55 | +        self.__process.start()
 | |
| 56 | + | |
| 57 | +        self.port = self.__queue.get()
 | |
| 58 | +        self.remote = 'localhost:{}'.format(self.port)
 | |
| 59 | + | |
| 60 | +    @staticmethod
 | |
| 61 | +    def serve(queue, instances, bot_session_queue, message_queue):
 | |
| 62 | +        pytest_cov.embed.cleanup_on_sigterm()
 | |
| 63 | + | |
| 64 | +        # Use max_workers default from Python 3.5+
 | |
| 65 | +        max_workers = (os.cpu_count() or 1) * 5
 | |
| 66 | +        server = grpc.server(futures.ThreadPoolExecutor(max_workers))
 | |
| 67 | +        port = server.add_insecure_port('localhost:0')
 | |
| 68 | + | |
| 69 | +        bots_service = service.BotsService(server)
 | |
| 70 | +        for name in instances:
 | |
| 71 | +            bots_interface = BotsInterface(bot_session_queue, message_queue)
 | |
| 72 | +            bots_service.add_instance(name, bots_interface)
 | |
| 73 | + | |
| 74 | +        server.start()
 | |
| 75 | +        queue.put(port)
 | |
| 76 | +        signal.pause()
 | |
| 77 | + | |
| 78 | +    def get_bot_session(self, timeout=1):
 | |
| 79 | +        bot_session = bots_pb2.BotSession()
 | |
| 80 | +        bot_session.ParseFromString(self.__bot_session_queue.get(timeout=timeout))
 | |
| 81 | +        return bot_session
 | |
| 82 | + | |
| 83 | +    # Injects leases
 | |
| 84 | +    def inject_work(self, lease=None, timeout=1):
 | |
| 85 | +        if not lease:
 | |
| 86 | +            lease = bots_pb2.Lease()
 | |
| 87 | +            lease.state = LeaseState.PENDING.value
 | |
| 88 | + | |
| 89 | +        lease_string = lease.SerializeToString()
 | |
| 90 | +        self.__message_queue.put(('INJECT_WORK', lease_string))
 | |
| 91 | + | |
| 92 | +    # Triggers a cancellation of a lease from server
 | |
| 93 | +    def cancel_lease(self, lease_id):
 | |
| 94 | +        self.__message_queue.put(('CANCEL_LEASE', lease_id))
 | |
| 95 | + | |
| 96 | +    def quit(self):
 | |
| 97 | +        if self.__process:
 | |
| 98 | +            self.__process.terminate()
 | |
| 99 | +            self.__process.join()
 | |
| 100 | + | |
| 101 | + | |
| 102 | +class BotsInterface:
 | |
| 103 | + | |
| 104 | +    def __init__(self, bot_session_queue, message_queue):
 | |
| 105 | +        self.__bot_session_queue = bot_session_queue
 | |
| 106 | +        self.__message_queue = message_queue
 | |
| 107 | + | |
| 108 | +    def register_instance_with_server(self, instance_name, server):
 | |
| 109 | +        server.add_bots_interface(self, instance_name)
 | |
| 110 | + | |
| 111 | +    def create_bot_session(self, parent, bot_session):
 | |
| 112 | +        name = "{}/{}".format(parent, str(uuid.uuid4()))
 | |
| 113 | +        bot_session.name = name
 | |
| 114 | + | |
| 115 | +        while not self.__message_queue.empty():
 | |
| 116 | +            message = self.__message_queue.get()
 | |
| 117 | +            if message[0] == 'INJECT_WORK':
 | |
| 118 | +                lease_string = message[1]
 | |
| 119 | +                lease = bots_pb2.Lease()
 | |
| 120 | +                lease.ParseFromString(lease_string)
 | |
| 121 | +                bot_session.leases.extend([lease])
 | |
| 122 | + | |
| 123 | +        self.__bot_session_queue.put(bot_session.SerializeToString())
 | |
| 124 | +        return bot_session
 | |
| 125 | + | |
| 126 | +    def update_bot_session(self, name, bot_session):
 | |
| 127 | +        for lease in bot_session.leases:
 | |
| 128 | +            state = LeaseState(lease.state)
 | |
| 129 | +            if state == LeaseState.COMPLETED:
 | |
| 130 | +                lease.Clear()
 | |
| 131 | + | |
| 132 | +            elif state == LeaseState.CANCELLED:
 | |
| 133 | +                lease.Clear()
 | |
| 134 | + | |
| 135 | +        while not self.__message_queue.empty():
 | |
| 136 | +            message = self.__message_queue.get()
 | |
| 137 | + | |
| 138 | +            if message[0] == 'INJECT_WORK':
 | |
| 139 | +                lease_string = message[1]
 | |
| 140 | +                lease = bots_pb2.Lease()
 | |
| 141 | +                lease.ParseFromString(lease_string)
 | |
| 142 | +                bot_session.leases.extend([lease])
 | |
| 143 | + | |
| 144 | +            elif message[0] == 'CANCEL_LEASE':
 | |
| 145 | +                lease_id = message[1]
 | |
| 146 | +                for lease in bot_session.leases:
 | |
| 147 | +                    if lease.id == lease_id:
 | |
| 148 | +                        lease.state = LeaseState.CANCELLED.value
 | |
| 149 | + | |
| 150 | +        self.__bot_session_queue.put(bot_session.SerializeToString())
 | |
| 151 | +        return bot_session | 
