Fix failing doctest