Fix more unit tests.

Finally fix #3 and #4 properly.