image-classification-service/test/unit_tests/mocked_server_test.py

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"