def test_place_with_resource_constraints(self, client_class):
        client_class.side_effect = self.create_fake_client

        bar_client = MagicMock()
        bar_response = PlaceResponse(PlaceResultCode.OK, agent_id="bar",
                                     score=Score(5, 90))
        bar_client.place.return_value = bar_response
        self._clients["bar"] = bar_client

        baz_client = MagicMock()
        baz_response = PlaceResponse(PlaceResultCode.OK, agent_id="baz",
                                     score=Score(30, 80))
        baz_client.place.return_value = baz_response
        self._clients["baz"] = baz_client

        scheduler = BranchScheduler("foo", 9)
        scheduler.configure([
            ChildInfo(id="bar", address="bar", constraints=[
                ResourceConstraint(ResourceConstraintType.DATASTORE, ["1"])]),
            ChildInfo(id="baz", address="baz", constraints=[
                ResourceConstraint(ResourceConstraintType.DATASTORE, ["2"])])])

        request = self._place_request()
        request.resource.vm.resource_constraints = [ResourceConstraint(
            ResourceConstraintType.DATASTORE, ["1"])]
        response = scheduler.place(request)
        assert_that(response.result, is_(PlaceResultCode.OK))
        assert_that(response.agent_id, is_("bar"))
    def test_place_sampling(self, shuffle, client_class):
        client_class.side_effect = self.create_fake_client
        shuffle.side_effect = self.fake_shuffle(
            [ChildInfo(id="bar", address="bar"),
             ChildInfo(id="baz", address="baz"),
             ChildInfo(id="qux", address="qux")])

        bar_client = MagicMock()
        bar_response = PlaceResponse(PlaceResultCode.OK, agent_id="bar",
                                     score=Score(5, 90))
        bar_client.place.return_value = bar_response
        self._clients["bar"] = bar_client

        baz_client = MagicMock()
        baz_response = PlaceResponse(PlaceResultCode.OK, agent_id="baz",
                                     score=Score(30, 80))
        baz_client.place.return_value = baz_response
        self._clients["baz"] = baz_client

        scheduler = BranchScheduler("foo", 9)
        scheduler.configure([ChildInfo(id="qux", address="qux"),
                             ChildInfo(id="bar", address="bar"),
                             ChildInfo(id="baz", address="baz")])

        response = scheduler.place(self._place_request())
        assert_that(response, is_(same_instance(baz_response)))
    def test_place_with_resource_constraints_no_match(self):
        scheduler = BranchScheduler("foo", 9)
        scheduler.configure([
            ChildInfo(id="bar", constraints=[
                ResourceConstraint(ResourceConstraintType.DATASTORE, ["1"])]),
            ChildInfo(id="baz", constraints=[
                ResourceConstraint(ResourceConstraintType.DATASTORE, ["2"])])])

        request = self._place_request()
        request.resource.vm.resource_constraints = [ResourceConstraint(
            "datastore", ["never_found"])]
        response = scheduler.place(request)
        assert_that(response.result, is_(PlaceResultCode.RESOURCE_CONSTRAINT))
    def test_find_resource_not_found(self, client_class):
        client_class.side_effect = self.create_fake_client

        bar_client = MagicMock()
        bar_response = FindResponse(FindResultCode.NOT_FOUND)
        bar_client.find.return_value = bar_response
        self._clients["bar"] = bar_client

        baz_client = MagicMock()
        baz_response = FindResponse(FindResultCode.NOT_FOUND)
        baz_client.find.return_value = baz_response
        self._clients["baz"] = baz_client

        scheduler = BranchScheduler("foo", 9)
        scheduler.configure([ChildInfo(id="bar", address="bar"),
                             ChildInfo(id="baz", address="baz")])

        response = scheduler.find(FindRequest())
        assert_that(response.result, is_(FindResultCode.NOT_FOUND))
    def test_place_timeout(self, client_class, wait_fn):
        client_class.side_effect = self.create_fake_client

        client = MagicMock()
        response = PlaceResponse(PlaceResultCode.OK)
        client.place.return_value = response
        self._clients["scheduler/baz"] = client

        wait_fn.return_value = set(), set()

        scheduler = BranchScheduler("foo", 9)
        scheduler.configure([ChildInfo(id="baz")])

        response = scheduler.place(self._place_request())
        assert_that(response.result, is_(PlaceResultCode.SYSTEM_ERROR))

        # must shutdown to join the futures since the patch lifecycle
        # ends when this method returns
        self._threadpool.shutdown()
    def test_place_resource_not_found(self, client_class):
        client_class.side_effect = self.create_fake_client

        bar_client = MagicMock()
        bar_response = PlaceResponse(PlaceResultCode.RESOURCE_CONSTRAINT)
        bar_client.place.return_value = bar_response
        self._clients["bar"] = bar_client

        baz_client = MagicMock()
        baz_response = PlaceResponse(PlaceResultCode.SYSTEM_ERROR)
        baz_client.place.return_value = baz_response
        self._clients["baz"] = baz_client

        scheduler = BranchScheduler("foo", 9)
        scheduler.configure([ChildInfo(id="bar", address="bar"),
                             ChildInfo(id="baz", address="baz")])

        response = scheduler.place(self._place_request())
        assert_that(response.result, is_(PlaceResultCode.RESOURCE_CONSTRAINT))
    def test_find_timeout(self, wait_fn, client_class):
        client_class.side_effect = self.create_fake_client

        client = MagicMock()
        response = FindResponse(FindResultCode.OK)
        client.find.return_value = response
        self._clients["baz"] = client

        scheduler = BranchScheduler("foo", 9)
        scheduler.configure([ChildInfo(id="baz", address="baz")])

        wait_fn.return_value = set(), set()

        response = scheduler.find(FindRequest())
        assert_that(response.result, is_(FindResultCode.NOT_FOUND))

        # must shutdown to join the futures since the patch lifecycle
        # ends when this method returns
        self._threadpool.shutdown()
    def test_find_sets_scheduler_id(self, client_class):
        client_class.side_effect = self.create_fake_client

        client = MagicMock()
        self._clients["bar"] = client

        scheduler = BranchScheduler("foo", 9)
        scheduler.configure([ChildInfo(id="bar", address="bar")])
        scheduler.find(FindRequest())

        client.find.assert_called_with(FindRequest(scheduler_id="bar"))