Fix tests importing six