refactor the tests for xx and return expected objects based on query