Update test_units.py

updated to reflect new WithUnit