Fix tests for new functionality.