Improve test coverage and integration tests

Edited by Gaba