Improved tests and docs