Fix unit testing bugs