60 lines
1.3 KiB
Python
60 lines
1.3 KiB
Python
import json
|
|
import logging
|
|
|
|
import pytest
|
|
|
|
from image_prediction.exceptions import IntentionalTestException
|
|
from image_prediction.flask import make_prediction_server
|
|
from image_prediction.utils import get_logger
|
|
|
|
|
|
@pytest.fixture
|
|
def mute_logger():
|
|
logger = get_logger()
|
|
level = logger.level
|
|
logger.setLevel(logging.CRITICAL + 1)
|
|
yield
|
|
logger.setLevel(level)
|
|
|
|
|
|
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("/ready")
|
|
assert response.status_code == 200
|
|
assert response.json == "OK"
|
|
|
|
|
|
def test_server_ready_check(client):
|
|
response = client.get("/health")
|
|
assert response.status_code == 200
|
|
assert response.json == "OK"
|