Improvement/tests

Refs #32 
Fixes #33