Brownhead/pytest

Added some unit testing.