Fixed tests to work with the functions after getting a new argument.