Add TODO tests