Cleaned up the test class a bit