Test improvements

Fixes #309