Refactor current state for tests to pass. Increase coverage