test_api.py

from multiprocessing import Process

import requests
import unittest

from conf.testing import czifile
from model_server.model import DummyImageToImageModel

class TestApiFromAutomatedClient(unittest.TestCase):
    def setUp(self) -> None:
        import uvicorn
        host = '127.0.0.1'
        port = 5000

        self.server_process = Process(
            target=uvicorn.run,
            args=('api:app', ),
            kwargs={'host': host, 'port': port, 'log_level': 'info'},
            daemon=True
        )
        self.uri = f'http://{host}:{port}/'
        self.server_process.start()

    @staticmethod
    def copy_input_file_to_server():
        import pathlib
        from shutil import copyfile
        from conf.server import paths

        outpath = pathlib.Path(paths['images']['inbound'] / czifile['filename'])

        copyfile(
            czifile['path'],
            outpath
        )

    def tearDown(self) -> None:
        self.server_process.terminate()

    def test_trivial_api_response(self):
        resp = requests.get(self.uri, )
        self.assertEqual(resp.status_code, 200)

    def test_list_empty_loaded_models(self):
        resp = requests.get(self.uri + 'models')
        self.assertEqual(resp.status_code, 200)
        self.assertEqual(resp.content, b'{}')

    def test_load_model(self):
        model_id = DummyImageToImageModel.model_id
        resp_load = requests.get(self.uri + f'models/{model_id}/load')
        self.assertEqual(resp_load.status_code, 200)
        resp_list = requests.get(self.uri + 'models')
        self.assertEqual(resp_list.status_code, 200)
        self.assertEqual(resp_list.content, b'{"dummy_make_white_square":"DummyImageToImageModel"}')

    def test_i2i_inference_errors_model_not_found(self):
        model_id = 'not_a_real_model'
        resp = requests.put(
            self.uri + f'i2i/infer/{model_id}',
            params={'input_filename': 'not_a_real_file.name'}
        )
        print(resp.content)
        self.assertEqual(resp.status_code, 409)

    def test_i2i_dummy_inference_by_api(self):
        model = DummyImageToImageModel()
        resp_load = requests.get(self.uri + f'models/{model.model_id}/load')
        self.assertEqual(resp_load.status_code, 200, f'Error loading {model.model_id}')
        self.copy_input_file_to_server()
        resp_infer = requests.put(
            self.uri + f'i2i/infer/{model.model_id}',
            params={'input_filename': czifile['filename']},
        )
        self.assertEqual(resp_infer.status_code, 200, f'Error inferring from {model.model_id}')