Add remainder of tests

I added the remainder of the tests, including some slight refactorings. I also added the failing test for the feature we need to add.