def test_correct(self):
        self.replace_corrupt_and_heavy_models()
        request = service_pb2.LoadModelsRequest(models=(
            model_pb2.ModelSpec(name="corrupt",
                                base_path="test/resources/models/corrupt"),
            model_pb2.ModelSpec(name="heavy",
                                base_path="test/resources/models/heavy"),
        ))
        response = self.stub.LoadModels(request)
        self.assertTrue(response.success)

        # Check both models' state is LOADED
        request = service_pb2.ModelStatusRequest(model=model_pb2.ModelSpec(
            name="corrupt"))
        response = self.stub.GetModelStatus(request)
        self.assertTrue(
            model_pb2.ModelStatus.ModelState.Name(response.status.state) ==
            "LOADED")
        request = service_pb2.ModelStatusRequest(model=model_pb2.ModelSpec(
            name="heavy"))
        response = self.stub.GetModelStatus(request)
        self.assertTrue(
            model_pb2.ModelStatus.ModelState.Name(response.status.state) ==
            "LOADED")
        self.revert_model_changes()
    def test_corrupt_model(self):
        model_spec = model_pb2.ModelSpec(
            name="corrupt", base_path="test/resources/models/corrupt")
        request = service_pb2.LoadModelsRequest(models=(model_spec, ))
        response = self.stub.LoadModels(request)
        self.assertFalse(response.success)

        # Check model's state is FAILED
        request = service_pb2.ModelStatusRequest(model=model_spec)
        response = self.stub.GetModelStatus(request)
        self.assertTrue(
            model_pb2.ModelStatus.ModelState.Name(response.status.state) ==
            "FAILED")
 def test_missing_model_name(self):
     request = service_pb2.LoadModelsRequest(models=(model_pb2.ModelSpec(
         base_path="test/resources/headers"), ))
     response = self.stub.LoadModels(request)
     self.assertFalse(response.success)
 def test_bad_model_path(self):
     request = service_pb2.LoadModelsRequest(models=(model_pb2.ModelSpec(
         name="foo", base_path="test/resources/models/not_existing"), ))
     response = self.stub.LoadModels(request)
     self.assertFalse(response.success)
 def test_empty_models(self):
     try:
         self.stub.LoadModels(service_pb2.LoadModelsRequest(models=()))
     except grpc.RpcError as error:
         self.assertTrue(
             error._state.code == grpc.StatusCode.INVALID_ARGUMENT)