Fix test by using mocked query