Fixing unit tests problem (#232)

Offending method was already put in a separate module, i just had to clean up the imports.