Add unit tests

A bit unnecessary, but can't hurt :)