test_workflow.py

import unittest

from conf.testing import czifile, output_path
from model_server.model import DummyImageToImageModel
from model_server.workflow import infer_image_to_image


class TestGetSessionObject(unittest.TestCase):
    def setUp(self) -> None:
        self.model = DummyImageToImageModel()

    def test_single_session_instance(self):
        result = infer_image_to_image(czifile['path'], self.model, output_path, channel=2)
        self.assertTrue(result.success)

        import tifffile
        img = tifffile.imread(result.output_filepath)
        w = czifile['w']
        h = czifile['h']

        self.assertEqual(
            img.shape,
            (h, w),
            'Inferred image is not the expected shape'
        )

        self.assertEqual(
            img[int(w/2), int(h/2)],
            255,
            'Middle pixel is not white as expected'
        )

        self.assertEqual(
            img[0, 0],
            0,
            'First pixel is not black as expected'
        )