fix(unit): Making unit tests support generator functions