fix missing import.
Merge pull request #1588 from nouiz/fix_test

fix missing import.