Added query method and updates unit tests