Update unit tests to new functionality provided
Update unit tests to new functionality provided