Query test improvements & minor refactoring