Fixing test coverage