Fix tests, add a new one too