import json import pytest from image_prediction.exceptions import IntentionalTestException from image_prediction.flask import make_prediction_server def predict_fn(x: bytes): x = int(x.decode()) if x == 42: return True else: raise IntentionalTestException("This is intended.") @pytest.fixture def server(): server = make_prediction_server(predict_fn) server.config.update({"TESTING": True}) return server @pytest.fixture def client(server): return server.test_client() def test_server_predict_success(client, mute_logger): response = client.post("/predict", data="42") assert json.loads(response.data) def test_server_predict_failure(client, mute_logger): response = client.post("/predict", data="13") assert response.status_code == 500 def test_server_health_check(client): response = client.get("/health") assert response.status_code == 200 assert response.json == "OK" def test_server_ready_check(client): response = client.get("/ready") assert response.status_code == 200 assert response.json == "OK"