Fixed tests to work with latest code.