Clean up test.
Merge pull request #1 from tonysyu/pr/613

Minor cleanups to test