Unit test fix

Steve, please check and merge. Fixed a bug yesterday's commit.