test: fix tests for #251

[xukai92]

addressing #254

i left comments in the code diff

[codecov[bot]]

Codecov Report

Patch has no changes to coverable lines.

:loudspeaker: Thoughts on this report? Let us know!.

[Roger-luo]

Thanks for the quick fix!