seveibar/nema17-stepper-motor-driver-tmc2209-usb-pd

firmware/test_tmc2209.py

import importlib.util
from pathlib import Path
import sys, types, unittest
from unittest.mock import patch

class ProtocolTests(unittest.TestCase):
    def setUp(self):
        self.now=0
        def ticks():
            self.now+=1
            return self.now
        with patch.dict(sys.modules, {'machine':types.SimpleNamespace(UART=None,Pin=None),
             'time':types.SimpleNamespace(sleep_ms=lambda _:None,ticks_ms=ticks,ticks_diff=lambda a,b:a-b)}):
            spec=importlib.util.spec_from_file_location('tmc_protocol',Path(__file__).with_name('tmc2209.py'))
            self.mod=importlib.util.module_from_spec(spec);spec.loader.exec_module(self.mod)
        mod=self.mod
        class Bus:
            def __init__(self):self.rx=b'';self.regs={6:0x21000000,2:0,1:0,0x6F:0};self.reject=False;self.corrupt=False
            def read(self):r,self.rx=self.rx,b'';return r
            def write(self,b):
                assert mod.crc8(b[:-1])==b[-1]
                if len(b)==8:
                    if not self.reject:
                        register=b[2]&127;value=int.from_bytes(b[3:7],'big')
                        self.regs[register]=self.regs.get(register,0)&~value if register==1 else value
                        self.regs[2]=(self.regs[2]+1)&255
                    self.rx=b
                else:
                    r=bytes((5,255,b[2]))+self.regs.get(b[2],0).to_bytes(4,'big')
                    self.rx=b+r+bytes((mod.crc8(r)^(1 if self.corrupt else 0),))
        self.bus=Bus();self.driver=mod.TMC2209(self.bus)
    def test_known_crc_vector(self):self.assertEqual(self.mod.crc8(bytes((5,0,0))),0x48)
    def test_echo_and_configuration(self):
        self.driver.configure();self.assertEqual(self.bus.regs[0],0x1C0)
        self.assertEqual((self.bus.regs[0x10]>>8)&31,14)
        self.assertEqual((self.bus.regs[0x6C]>>24)&15,4)
        self.assertEqual(self.bus.regs[0x22],0)
    def test_rejected_write_and_wrap(self):
        self.bus.regs[2]=255;self.driver.write(0,123);self.assertEqual(self.bus.regs[2],0)
        self.bus.reject=True
        with self.assertRaises(OSError):self.driver.write(0,0)
    def test_bad_crc(self):
        self.bus.corrupt=True
        with self.assertRaises(OSError):self.driver.read(6)
    def test_reset_thermal_short_and_wrong_part(self):
        self.driver.configure();self.bus.regs[1]=1
        with self.assertRaises(OSError):self.driver.check()
        self.driver.configure();self.bus.regs[0x6F]=1
        with self.assertRaises(OSError):self.driver.check()
        self.bus.regs[6]=0
        with self.assertRaises(OSError):self.driver.configure()
if __name__=='__main__':unittest.main()