add import statements to test.py
fix docstring for wem module