Wrote tests

* Wrote unit tests to verify functionality
* Modified classes to be more easily testable