Adding unit tests for new functionality.