broken tests.py
Fix tests.py
Add unit tests