Test Generation fix