Minor test edits

Cleaning up code.