Test cleanups and improvements