WIP: add tests