Added correct import for behave tests (to fix the bug)