Revert "Revert "Fixed unittest speed""

This reverts commit fd32e197e1d618773bf425642f3887fc20fffa29.