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.put(
            self.uri + f'models/load',
            params={'model_id': model_id}
        )
        self.assertEqual(resp_load.status_code, 200)
        resp_list = requests.get(self.uri + 'models')
        self.assertEqual(resp_list.status_code, 200)
        rj = resp_list.json()
        self.assertEqual(rj[model_id]['class'], 'DummyImageToImageModel')


    def test_i2i_inference_errors_model_not_found(self):
        model_id = 'not_a_real_model'
        resp = requests.put(
            self.uri + f'i2i/infer/',
            params={
                'model_id': model_id,
                '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.put(
            self.uri + f'models/load',
            params={'model_id': model.model_id}
        )
        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/',
            params={
                'model_id': model.model_id,
                'input_filename': czifile['filename'],
                'channel': 2,
            },
        )
        self.assertEqual(resp_infer.status_code, 200, f'Error inferring from {model.model_id}')