gen_esp32part.py: Remaining Python 3 changes for unit tests to pass, plus unit tests

Ref https://github.com/espressif/esp-idf/pull/577
This commit is contained in:
Angus Gratton 2017-05-12 12:25:41 +10:00 committed by Angus Gratton
parent beffcd6468
commit aaa8170865
2 changed files with 61 additions and 52 deletions

View File

@ -6,6 +6,7 @@
# #
# See http://esp-idf.readthedocs.io/en/latest/partition-tables.html for explanation of # See http://esp-idf.readthedocs.io/en/latest/partition-tables.html for explanation of
# partition table structure and uses. # partition table structure and uses.
from __future__ import print_function, division
import argparse import argparse
import os import os
import re import re
@ -94,7 +95,7 @@ class PartitionTable(list):
data = b[o:o+32] data = b[o:o+32]
if len(data) != 32: if len(data) != 32:
raise InputError("Partition table length must be a multiple of 32 bytes") raise InputError("Partition table length must be a multiple of 32 bytes")
if data == '\xFF'*32: if data == b'\xFF'*32:
return result # got end marker return result # got end marker
result.append(PartitionDefinition.from_binary(data)) result.append(PartitionDefinition.from_binary(data))
raise InputError("Partition table is missing an end-of-table marker") raise InputError("Partition table is missing an end-of-table marker")
@ -246,8 +247,9 @@ class PartitionDefinition(object):
res = cls() res = cls()
(magic, res.type, res.subtype, res.offset, (magic, res.type, res.subtype, res.offset,
res.size, res.name, flags) = struct.unpack(cls.STRUCT_FORMAT, b) res.size, res.name, flags) = struct.unpack(cls.STRUCT_FORMAT, b)
if "\x00" in res.name: # strip null byte padding from name string if b"\x00" in res.name: # strip null byte padding from name string
res.name = res.name[:res.name.index("\x00")] res.name = res.name[:res.name.index(b"\x00")]
res.name = res.name.decode()
if magic != cls.MAGIC_BYTES: if magic != cls.MAGIC_BYTES:
raise InputError("Invalid magic bytes (%r) for partition definition" % magic) raise InputError("Invalid magic bytes (%r) for partition definition" % magic)
for flag,bit in cls.FLAGS.items(): for flag,bit in cls.FLAGS.items():
@ -275,7 +277,7 @@ class PartitionDefinition(object):
if not simple_formatting and include_sizes: if not simple_formatting and include_sizes:
for (val, suffix) in [ (0x100000, "M"), (0x400, "K") ]: for (val, suffix) in [ (0x100000, "M"), (0x400, "K") ]:
if a % val == 0: if a % val == 0:
return "%d%s" % (a / val, suffix) return "%d%s" % (a // val, suffix)
return "0x%x" % a return "0x%x" % a
def lookup_keyword(t, keywords): def lookup_keyword(t, keywords):
@ -323,7 +325,7 @@ def main():
parser.add_argument('--verify', '-v', help='Verify partition table fields', default=True, action='store_false') parser.add_argument('--verify', '-v', help='Verify partition table fields', default=True, action='store_false')
parser.add_argument('--quiet', '-q', help="Don't print status messages to stderr", action='store_true') parser.add_argument('--quiet', '-q', help="Don't print status messages to stderr", action='store_true')
parser.add_argument('input', help='Path to CSV or binary file to parse. Will use stdin if omitted.', type=argparse.FileType('r'), default=sys.stdin) parser.add_argument('input', help='Path to CSV or binary file to parse. Will use stdin if omitted.', type=argparse.FileType('rb'), default=sys.stdin)
parser.add_argument('output', help='Path to output converted binary or CSV file. Will use stdout if omitted, unless the --display argument is also passed (in which case only the summary is printed.)', parser.add_argument('output', help='Path to output converted binary or CSV file. Will use stdout if omitted, unless the --display argument is also passed (in which case only the summary is printed.)',
nargs='?', nargs='?',
default='-') default='-')
@ -337,6 +339,7 @@ def main():
status("Parsing binary partition input...") status("Parsing binary partition input...")
table = PartitionTable.from_binary(input) table = PartitionTable.from_binary(input)
else: else:
input = input.decode()
status("Parsing CSV input...") status("Parsing CSV input...")
table = PartitionTable.from_csv(input) table = PartitionTable.from_csv(input)
@ -357,5 +360,5 @@ if __name__ == '__main__':
try: try:
main() main()
except InputError as e: except InputError as e:
print >>sys.stderr, e print(e, file=sys.stderr)
sys.exit(2) sys.exit(2)

View File

@ -1,4 +1,5 @@
#!/usr/bin/env python #!/usr/bin/env python
from __future__ import print_function, division
import unittest import unittest
import struct import struct
import csv import csv
@ -14,38 +15,40 @@ SIMPLE_CSV = """
factory,0,2,65536,1048576, factory,0,2,65536,1048576,
""" """
LONGER_BINARY_TABLE = "" LONGER_BINARY_TABLE = b""
# type 0x00, subtype 0x00, # type 0x00, subtype 0x00,
# offset 64KB, size 1MB # offset 64KB, size 1MB
LONGER_BINARY_TABLE += "\xAA\x50\x00\x00" + \ LONGER_BINARY_TABLE += b"\xAA\x50\x00\x00" + \
"\x00\x00\x01\x00" + \ b"\x00\x00\x01\x00" + \
"\x00\x00\x10\x00" + \ b"\x00\x00\x10\x00" + \
"factory\0" + ("\0"*8) + \ b"factory\0" + (b"\0"*8) + \
"\x00\x00\x00\x00" b"\x00\x00\x00\x00"
# type 0x01, subtype 0x20, # type 0x01, subtype 0x20,
# offset 0x110000, size 128KB # offset 0x110000, size 128KB
LONGER_BINARY_TABLE += "\xAA\x50\x01\x20" + \ LONGER_BINARY_TABLE += b"\xAA\x50\x01\x20" + \
"\x00\x00\x11\x00" + \ b"\x00\x00\x11\x00" + \
"\x00\x02\x00\x00" + \ b"\x00\x02\x00\x00" + \
"data" + ("\0"*12) + \ b"data" + (b"\0"*12) + \
"\x00\x00\x00\x00" b"\x00\x00\x00\x00"
# type 0x10, subtype 0x00, # type 0x10, subtype 0x00,
# offset 0x150000, size 1MB # offset 0x150000, size 1MB
LONGER_BINARY_TABLE += "\xAA\x50\x10\x00" + \ LONGER_BINARY_TABLE += b"\xAA\x50\x10\x00" + \
"\x00\x00\x15\x00" + \ b"\x00\x00\x15\x00" + \
"\x00\x10\x00\x00" + \ b"\x00\x10\x00\x00" + \
"second" + ("\0"*10) + \ b"second" + (b"\0"*10) + \
"\x00\x00\x00\x00" b"\x00\x00\x00\x00"
LONGER_BINARY_TABLE += "\xFF" * 32 LONGER_BINARY_TABLE += b"\xFF" * 32
def _strip_trailing_ffs(binary_table): def _strip_trailing_ffs(binary_table):
""" """
Strip all FFs down to the last 32 bytes (terminating entry) Strip all FFs down to the last 32 bytes (terminating entry)
""" """
while binary_table.endswith("\xFF"*64): while binary_table.endswith(b"\xFF"*64):
binary_table = binary_table[0:len(binary_table)-32] binary_table = binary_table[0:len(binary_table)-32]
return binary_table return binary_table
class CSVParserTests(unittest.TestCase): class CSVParserTests(unittest.TestCase):
def test_simple_partition(self): def test_simple_partition(self):
@ -166,8 +169,8 @@ first, 0x30, 0xEE, 0x100400, 0x300000
t = PartitionTable.from_csv(csv) t = PartitionTable.from_csv(csv)
tb = _strip_trailing_ffs(t.to_binary()) tb = _strip_trailing_ffs(t.to_binary())
self.assertEqual(len(tb), 64) self.assertEqual(len(tb), 64)
self.assertEqual('\xAA\x50', tb[0:2]) # magic self.assertEqual(b'\xAA\x50', tb[0:2]) # magic
self.assertEqual('\x30\xee', tb[2:4]) # type, subtype self.assertEqual(b'\x30\xee', tb[2:4]) # type, subtype
eo, es = struct.unpack("<LL", tb[4:12]) eo, es = struct.unpack("<LL", tb[4:12])
self.assertEqual(eo, 0x100400) # offset self.assertEqual(eo, 0x100400) # offset
self.assertEqual(es, 0x300000) # size self.assertEqual(es, 0x300000) # size
@ -180,8 +183,8 @@ second,0x31, 0xEF, , 0x100000
t = PartitionTable.from_csv(csv) t = PartitionTable.from_csv(csv)
tb = _strip_trailing_ffs(t.to_binary()) tb = _strip_trailing_ffs(t.to_binary())
self.assertEqual(len(tb), 96) self.assertEqual(len(tb), 96)
self.assertEqual('\xAA\x50', tb[0:2]) self.assertEqual(b'\xAA\x50', tb[0:2])
self.assertEqual('\xAA\x50', tb[32:34]) self.assertEqual(b'\xAA\x50', tb[32:34])
def test_encrypted_flag(self): def test_encrypted_flag(self):
@ -200,12 +203,12 @@ class BinaryParserTests(unittest.TestCase):
def test_parse_one_entry(self): def test_parse_one_entry(self):
# type 0x30, subtype 0xee, # type 0x30, subtype 0xee,
# offset 1MB, size 2MB # offset 1MB, size 2MB
entry = "\xAA\x50\x30\xee" + \ entry = b"\xAA\x50\x30\xee" + \
"\x00\x00\x10\x00" + \ b"\x00\x00\x10\x00" + \
"\x00\x00\x20\x00" + \ b"\x00\x00\x20\x00" + \
"0123456789abc\0\0\0" + \ b"0123456789abc\0\0\0" + \
"\x00\x00\x00\x00" + \ b"\x00\x00\x00\x00" + \
"\xFF" * 32 b"\xFF" * 32
# verify that parsing 32 bytes as a table # verify that parsing 32 bytes as a table
# or as a single Definition are the same thing # or as a single Definition are the same thing
t = PartitionTable.from_binary(entry) t = PartitionTable.from_binary(entry)
@ -240,33 +243,36 @@ class BinaryParserTests(unittest.TestCase):
self.assertEqual(round_trip, LONGER_BINARY_TABLE) self.assertEqual(round_trip, LONGER_BINARY_TABLE)
def test_bad_magic(self): def test_bad_magic(self):
bad_magic = "OHAI" + \ bad_magic = b"OHAI" + \
"\x00\x00\x10\x00" + \ b"\x00\x00\x10\x00" + \
"\x00\x00\x20\x00" + \ b"\x00\x00\x20\x00" + \
"0123456789abc\0\0\0" + \ b"0123456789abc\0\0\0" + \
"\x00\x00\x00\x00" b"\x00\x00\x00\x00"
with self.assertRaisesRegexp(InputError, "Invalid magic bytes"): with self.assertRaisesRegexp(InputError, "Invalid magic bytes"):
PartitionTable.from_binary(bad_magic) PartitionTable.from_binary(bad_magic)
def test_bad_length(self): def test_bad_length(self):
bad_length = "OHAI" + \ bad_length = b"OHAI" + \
"\x00\x00\x10\x00" + \ b"\x00\x00\x10\x00" + \
"\x00\x00\x20\x00" + \ b"\x00\x00\x20\x00" + \
"0123456789" b"0123456789"
with self.assertRaisesRegexp(InputError, "32 bytes"): with self.assertRaisesRegexp(InputError, "32 bytes"):
PartitionTable.from_binary(bad_length) PartitionTable.from_binary(bad_length)
class CSVOutputTests(unittest.TestCase): class CSVOutputTests(unittest.TestCase):
def _readcsv(self, source_str):
return list(csv.reader(source_str.split("\n")))
def test_output_simple_formatting(self): def test_output_simple_formatting(self):
table = PartitionTable.from_csv(SIMPLE_CSV) table = PartitionTable.from_csv(SIMPLE_CSV)
as_csv = table.to_csv(True) as_csv = table.to_csv(True)
c = csv.reader(as_csv.split("\n")) c = self._readcsv(as_csv)
# first two lines should start with comments # first two lines should start with comments
self.assertEqual(c.next()[0][0], "#") self.assertEqual(c[0][0][0], "#")
self.assertEqual(c.next()[0][0], "#") self.assertEqual(c[1][0][0], "#")
row = c.next() row = c[2]
self.assertEqual(row[0], "factory") self.assertEqual(row[0], "factory")
self.assertEqual(row[1], "0") self.assertEqual(row[1], "0")
self.assertEqual(row[2], "2") self.assertEqual(row[2], "2")
@ -280,11 +286,11 @@ class CSVOutputTests(unittest.TestCase):
def test_output_smart_formatting(self): def test_output_smart_formatting(self):
table = PartitionTable.from_csv(SIMPLE_CSV) table = PartitionTable.from_csv(SIMPLE_CSV)
as_csv = table.to_csv(False) as_csv = table.to_csv(False)
c = csv.reader(as_csv.split("\n")) c = self._readcsv(as_csv)
# first two lines should start with comments # first two lines should start with comments
self.assertEqual(c.next()[0][0], "#") self.assertEqual(c[0][0][0], "#")
self.assertEqual(c.next()[0][0], "#") self.assertEqual(c[1][0][0], "#")
row = c.next() row = c[2]
self.assertEqual(row[0], "factory") self.assertEqual(row[0], "factory")
self.assertEqual(row[1], "app") self.assertEqual(row[1], "app")
self.assertEqual(row[2], "2") self.assertEqual(row[2], "2")
@ -303,7 +309,7 @@ class CommandLineTests(unittest.TestCase):
csvpath = tempfile.mktemp() csvpath = tempfile.mktemp()
# copy binary contents to temp file # copy binary contents to temp file
with open(binpath, 'w') as f: with open(binpath, 'wb') as f:
f.write(LONGER_BINARY_TABLE) f.write(LONGER_BINARY_TABLE)
# run gen_esp32part.py to convert binary file to CSV # run gen_esp32part.py to convert binary file to CSV