Fix EntityStore usage in tests

Change-Id: I010755f2cfd7be2eedb4dc83cba56a3c47b702a5