Fix unit tests + add more unit tests.