Update with unit tests