Some cleanup of QueryEntity and its test

Change-Id: I1d9edaeb96ba8f5dd32d8504a2dfda5dbe696353