fixing bugs, restore unit tests