Update tests to match new generator output