Revert "Revert "unit test fix""

This reverts commit 258d9c627fbfb8523dd1d79d4ad8baa76108e648.