Fix unit tests
Merge branch 'revamp'