Fix the test code.

Previous commit add an extra argunment to the "pre_generation" function which I forgot to update the test code.

It's all green now.