Add code coverage, improve unit tests