test_api.py

from multiprocessing import Process
import requests
import unittest

from conf.testing import czifile, output_path
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()

    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 = requests.get(self.uri + f'models/{model_id}/load')
        self.assertEqual(resp.status_code, 200)
        loaded = requests.get(self.uri + 'models')
        self.assertEqual(loaded.content, b'{"model_id":"DummyImageToImageModel"}')

    def test_i2i_inference_errors_model_not_sound(self):
        model_id = 'not_a_real_model'
        resp = requests.post(self.uri + f'i2i/infer/{model_id}')
        self.assertEqual(resp.status_code, 404)

    def test_i2i_dummy_inference_by_api(self):
        model = DummyImageToImageModel()
        model_id = model.model_id
        resp = requests.post(
            self.uri + f'/i2i/infer/{model_id}',
            str(czifile['path']),
        )
        print(resp)
        self.assertEqual(resp.status_code, 200)