Fix tests to pass with the recent changes.

Also make a small style change in the main function. And move the noop
main block for unit tests outside of the function body definition (this
is simply for the compiler not to complain about unreachable statements
on the `unittest` version).