Fix unit tests
Cleanup unit tests