Fixed up our test
Cleaned up code