Fixing up code in test so that it compiles.

Library has been refactored recently, but tests have not been updated,
fixing that.