Fixed all the tests broken due to refactor

All tests are running on local. Hope to set up things somehow on
travis.