Added some unit tests; plus minor changes