Cleaned up a bit. Now all tests pass.