Tighten testing logic a bit