fixed bug and updating tests