add tests, fix review notes