Fix test.
Merge pull request #79 from eric-wu/upgrade

Fix test