Updates and small unit test