from types import SimpleNamespace
from unittest import TestCase
import numpy as np
from station.image_preview import jpeg_preview


class ImagePreviewTest(TestCase):
    def test_padding_vector_and_invalid_buffers(self):
        rgb=SimpleNamespace(encoding='rgb8',width=2,height=2,step=8,is_bigendian=False,data=(bytes([90,10,30])*2+b'xx')*2)
        self.assertTrue(jpeg_preview(rgb).startswith(b'\xff\xd8'))
        field=np.ones((4,4,2),np.float32)
        msg=SimpleNamespace(encoding='32FC2',width=4,height=4,step=32,is_bigendian=False,data=field.tobytes())
        self.assertTrue(jpeg_preview(msg).startswith(b'\xff\xd8'))
        field[0,0,0]=np.nan;msg.data=field.tobytes()
        self.assertIsNone(jpeg_preview(msg))
        rgb.data=b'bad';self.assertIsNone(jpeg_preview(rgb))