Fix tests and documentation