mirror of
https://github.com/ITotalJustice/sphaira.git
synced 2025-11-02 01:16:03 +01:00
115 lines
3.9 KiB
Python
115 lines
3.9 KiB
Python
import unittest
|
|
import crc32c
|
|
import os
|
|
|
|
from usb_common import CMD_EXPORT, CMD_QUIT, RESULT_OK, RESULT_ERROR
|
|
|
|
class FakeUsb:
|
|
def __init__(self, files=None):
|
|
# files: list of tuples (filename: str, data: bytes)
|
|
self.files = files or [("testfile.bin", b"testdata")]
|
|
self._cmd_index = 0
|
|
self._file_index = 0
|
|
self._data_index = 0
|
|
self.results = []
|
|
self._reading_filename = True
|
|
self._reading_data = False
|
|
self._current_data = b""
|
|
self._current_data_offset = 0
|
|
self._current_data_sent = 0
|
|
self._current_file = None
|
|
self._send_data_header_calls = 0
|
|
|
|
def wait_for_connect(self):
|
|
pass
|
|
|
|
def get_send_header(self):
|
|
# Simulate command sequence: export for each file, then quit
|
|
if self._cmd_index < len(self.files):
|
|
filename, data = self.files[self._cmd_index]
|
|
self._current_file = (filename, data)
|
|
self._cmd_index += 1
|
|
self._reading_filename = True
|
|
self._reading_data = False
|
|
self._current_data = data
|
|
self._current_data_offset = 0
|
|
self._current_data_sent = 0
|
|
self._send_data_header_calls = 0
|
|
return [CMD_EXPORT, len(filename.encode("utf-8")), 0]
|
|
else:
|
|
return [CMD_QUIT, 0, 0]
|
|
|
|
def read(self, size):
|
|
# Simulate reading file name or data
|
|
if self._reading_filename:
|
|
filename = self._current_file[0].encode("utf-8")
|
|
self._reading_filename = False
|
|
self._reading_data = True
|
|
return filename[:size]
|
|
elif self._reading_data:
|
|
# Return file data for export
|
|
data = self._current_data[self._current_data_sent:self._current_data_sent+size]
|
|
self._current_data_sent += len(data)
|
|
return data
|
|
else:
|
|
return b""
|
|
|
|
def get_send_data_header(self):
|
|
# Simulate sending data in one chunk, then finish
|
|
if self._send_data_header_calls == 0:
|
|
self._send_data_header_calls += 1
|
|
data = self._current_data
|
|
crc = crc32c.crc32c(data)
|
|
return [0, len(data), crc]
|
|
else:
|
|
return [0, 0, 0] # End of transfer
|
|
|
|
def send_result(self, result):
|
|
self.results.append(result)
|
|
|
|
# test case for usb_export.py
|
|
class TestUsbExport(unittest.TestCase):
|
|
def setUp(self):
|
|
self.root = "test_output"
|
|
os.makedirs(self.root, exist_ok=True)
|
|
# 100 files named test1.bin, test2.bin, ..., test10.bin, each with different sizes
|
|
self.files = [
|
|
(f"test{i+1}.bin", bytes([65 + i]) * (i * 100 + 1)) for i in range(100)
|
|
]
|
|
self.fake_usb = FakeUsb(files=self.files)
|
|
|
|
def tearDown(self):
|
|
# Clean up created files/folders
|
|
for f in os.listdir(self.root):
|
|
os.remove(os.path.join(self.root, f))
|
|
os.rmdir(self.root)
|
|
|
|
def test_export_multiple_files(self):
|
|
from usb_export import get_file_name, create_file_folder, wait_for_input
|
|
|
|
# Simulate the main loop for all files
|
|
for filename, data in self.files:
|
|
cmd, name_len, _ = self.fake_usb.get_send_header()
|
|
self.assertEqual(cmd, CMD_EXPORT)
|
|
|
|
file_name = get_file_name(self.fake_usb, name_len)
|
|
self.assertEqual(file_name, filename)
|
|
|
|
full_path = create_file_folder(self.root, file_name)
|
|
self.fake_usb.send_result(RESULT_OK)
|
|
|
|
wait_for_input(self.fake_usb, full_path)
|
|
|
|
# Check file was created and contents match
|
|
with open(full_path, "rb") as f:
|
|
filedata = f.read()
|
|
self.assertEqual(filedata, data)
|
|
|
|
# After all files, should get CMD_QUIT
|
|
cmd, _, _ = self.fake_usb.get_send_header()
|
|
self.assertEqual(cmd, CMD_QUIT)
|
|
self.assertIn(RESULT_OK, self.fake_usb.results)
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|