Write some automated tests

[harrislapiroff]

It would be good to have some unit tests to speed up reviewing PRs in this repo. Probably using pytest or even just python's built in unit testing.