Tests: Code clean up