Better setup for tests
Fix bug
Misc changes