Improve test coverage

[10Electra]

Test coverage is only 60% in our model, it needs to be at least 80%