Cleaned up test code a bit