Fix failing test

Rearranged a little code in CLI module so that test would
work properly again.