MustafaMulla29/rp2040-bldc-motor-controller-new

firmware/test_pd_profile.py

import unittest
from pd_profile import PDOS, nvm_9v_image, nvm_is_9v, configure_runtime, verify_runtime
from pd_nvm import NVM

# STUSB4500 factory NVM example published by SparkFun. Reserved/calibration
# fields differ between revisions: the programming code preserves their bytes.
FACTORY = bytes.fromhex(
    '0000b0aa00450000 10409c1cff013cdf 02400f003200fcf1 '
    '001956aff5355f00 004b9021430040fb')

class Clock:
    now = 0
    def ticks_ms(self): return self.now
    def ticks_diff(self, a, b): return a-b
    def sleep_ms(self, delay): self.now += delay

class Device:
    """Bus-level model: sector erase/load/program plus NVM reload on reset."""
    def __init__(self):
        self.image = bytearray(FACTORY)
        self.regs = bytearray(256)
        self.programmed = []
        self.erase_mask = 0
        self.locked = True
        self.stuck = False
        self.corrupt = False
        self.reload()

    def reload(self):
        count = (self.image[26] >> 1) & 3
        volts = (100, (self.image[33] << 2) | (self.image[32] >> 6),
                 ((self.image[35] & 3) << 8) | self.image[34])
        codes = (self.image[26] >> 4, self.image[28] & 15, self.image[29] >> 4)
        self.regs[0x70] = count
        for i, (voltage, code) in enumerate(zip(volts, codes)):
            amps = 25*(code+1) if code <= 11 else 50*code-250
            pdo = (voltage << 10) | amps
            if i == 0 and self.image[26] & 1: pdo |= 1 << 26
            self.regs[0x85+4*i:0x89+4*i] = pdo.to_bytes(4,'little')

    def readfrom_mem(self, address, reg, size):
        return bytes(self.regs[reg:reg+size])

    def writeto_mem(self, address, reg, data):
        self.regs[reg:reg+len(data)] = data
        if reg == 0x95: self.locked = data[0] != 0x47
        if reg == 0x96 and data[0] & 0x10:
            if self.locked: raise OSError('NVM locked')
            op = self.regs[0x97] & 7
            sector = data[0] & 7
            if op == 0:
                self.regs[0x53:0x5B] = self.image[8*sector:8*sector+8]
            elif op == 2: self.erase_mask = self.regs[0x97] >> 3
            elif op == 5:
                for s in range(5):
                    if self.erase_mask & (1 << s): self.image[8*s:8*s+8] = b'\xff'*8
            elif op == 1: self.loaded = bytes(self.regs[0x53:0x5B])
            elif op == 6:
                self.image[8*sector:8*sector+8] = self.loaded
                self.programmed.append(sector)
                if self.corrupt: self.image[8*sector] ^= 1
            if not self.stuck: self.regs[0x96] &= ~0x10

class PDProfileTests(unittest.TestCase):
    def test_factory_to_9v_preserves_unrelated_fields(self):
        desired = nvm_9v_image(FACTORY)
        self.assertFalse(nvm_is_9v(FACTORY))
        self.assertTrue(nvm_is_9v(desired))
        self.assertEqual(desired[:24], FACTORY[:24])
        self.assertEqual(desired[26], 0x15)
        self.assertEqual(desired[36] & 0x60, 0x40)
        self.assertEqual(desired[38] & 0x18, 0x10)
        self.assertEqual(desired[28:30], b'\xf7\x75')
        self.assertEqual(desired[32:36], b'\x00\x2d\xb4\x20')
        self.assertEqual(nvm_9v_image(desired), desired)
        for byte, mask in {26:0xff,28:0x0f,29:0xf0,32:0xc0,33:0xff,34:0xff,35:0x03,36:0x60,38:0x18}.items():
            self.assertEqual((desired[byte] ^ FACTORY[byte]) & ~mask, 0)
        changed = {i for i in range(40) if desired[i] != FACTORY[i]}
        self.assertTrue(changed <= {26,28,29,32,33,34,35,36,38})
        with self.assertRaises(ValueError): nvm_9v_image(b'\x00'*39)

    def test_partial_program_and_cold_start_without_runtime_writes(self):
        bus = Device()
        nvm = NVM(bus, clock=Clock())
        self.assertTrue(nvm.provision_9v())
        self.assertEqual(bus.programmed, [3,4])
        self.assertEqual(bus.erase_mask, 0x18)
        self.assertTrue(bus.locked)
        self.assertEqual(bus.image[:24], FACTORY[:24])
        bus.reload()  # only NVM load: no runtime profile setup to hide bad NVM
        verify_runtime(bus)
        self.assertFalse(nvm.provision_9v())
        self.assertEqual(bus.programmed, [3,4])

    def test_factory_runtime_profiles_and_corruption_rejected(self):
        bus = Device()
        with self.assertRaises(OSError): verify_runtime(bus)
        configure_runtime(bus)
        verify_runtime(bus)
        for pdo in PDOS: self.assertIn(((pdo >> 10) & 1023)*50, (5000,9000))
        for register, voltage in ((0x89, 300), (0x8D, 400)):
            # Reject 15 V active PDO2 and 20 V even in inactive PDO3.
            bus.regs[register:register+4] = ((voltage << 10)|200).to_bytes(4,'little')
            with self.assertRaises(OSError): verify_runtime(bus)
            configure_runtime(bus)
        for count in (0, 1, 3, 4, 6, 7):
            bus.regs[0x70] = count
            with self.assertRaises(OSError): verify_runtime(bus)
        configure_runtime(bus)
        self.assertEqual(bus.regs[0x70] & 7, 2)

    def test_timeout_and_corrupt_write_release_test_mode(self):
        bus = Device(); bus.stuck = True
        clock = Clock()
        with self.assertRaises(OSError): NVM(bus,clock=clock).read()
        self.assertEqual(clock.now,1000)
        self.assertTrue(bus.locked)
        bus = Device(); bus.corrupt = True
        with self.assertRaises(OSError): NVM(bus,clock=Clock()).provision_9v()
        self.assertTrue(bus.locked)

if __name__ == '__main__': unittest.main()