Fix tests

Fixes couple of tests that are failing in branches

This is own pull request since I want to run Travis tests against this.