Small changes in code. First unit test passes.