Fix tests

Adapted from @beatgammit's PR, with slight changes.

Closes #7.