Improve test coverage

Second PR on adding tests. I also did some refactoring to improve testability and corrected minor errors in argument parsing.