You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
MAI/tests/test_delimited_frame_captur...

91 lines
2.9 KiB
Python

import socket
import tempfile
import unittest
from pathlib import Path
from delimited_frame_capture import DelimitedFrameCapture
class DelimitedFrameCaptureTests(unittest.TestCase):
WIDTH = 4
HEIGHT = 2
SEPARATOR = 255
def frame_bytes(self, offset=0):
size = self.WIDTH * self.HEIGHT * 3
return bytes((offset + index) % 200 for index in range(size))
def test_reads_raw_frames_from_delimited_log(self):
first = self.frame_bytes()
second = self.frame_bytes(20)
with tempfile.TemporaryDirectory() as tmp:
path = Path(tmp) / "frames.dump"
path.write_bytes(first + bytes((self.SEPARATOR,)) + second)
cap = DelimitedFrameCapture(
source=path,
separator=self.SEPARATOR,
encoding="bgr24",
width=self.WIDTH,
height=self.HEIGHT,
)
try:
ok1, frame1 = cap.read()
ok2, frame2 = cap.read()
self.assertTrue(ok1)
self.assertTrue(ok2)
self.assertEqual(frame1.shape, (self.HEIGHT, self.WIDTH, 3))
self.assertEqual(int(frame2[0, 0, 0]), 20)
finally:
cap.release()
def test_reads_raw_frame_from_live_udp_stream(self):
frame_data = self.frame_bytes()
cap = DelimitedFrameCapture(
host="127.0.0.1",
port=0,
separator=self.SEPARATOR,
encoding="bgr24",
width=self.WIDTH,
height=self.HEIGHT,
)
sender = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
try:
sender.sendto(
frame_data + bytes((self.SEPARATOR,)),
("127.0.0.1", cap.port),
)
ok, frame = cap.read()
self.assertTrue(ok)
self.assertEqual(frame.shape, (self.HEIGHT, self.WIDTH, 3))
finally:
sender.close()
cap.release()
def test_live_raw_frame_may_contain_separator_byte(self):
frame_data = bytes(range(self.WIDTH * self.HEIGHT))
cap = DelimitedFrameCapture(
host="127.0.0.1",
port=0,
separator=0,
encoding="gray8",
width=self.WIDTH,
height=self.HEIGHT,
)
sender = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
try:
sender.sendto(frame_data[:3], ("127.0.0.1", cap.port))
sender.sendto(frame_data[3:], ("127.0.0.1", cap.port))
sender.sendto(b"\x00", ("127.0.0.1", cap.port))
ok, frame = cap.read()
self.assertTrue(ok)
self.assertEqual(frame.shape, (self.HEIGHT, self.WIDTH, 3))
self.assertEqual(int(frame[0, 0, 0]), 0)
self.assertEqual(int(frame[-1, -1, 0]), 7)
finally:
sender.close()
cap.release()
if __name__ == "__main__":
unittest.main()