Adding some simple functionality tests and fixing a bug