Small fix + update unit tests.