Fix for the fix

Forgot to copy the new class from the test file.