Add missing import

Also add tests

Fixes #511
Merge remote-tracking branch 'zyx-i/fix-511' into develop