Refactor integration tests
Finish cleaning up unit tests