Add unit test for query operation