Updating tests/code to latest changes