Added test
Minor fixes, now passes tests.