Fix test

Signed-off-by: Edward Cheng <c4150221@gmail.com>