More unit tests. A few bugs fixed. Solves part of #6