Minor MR with fixes and adding missing tests.
This will create a new commit in order to revert the existing changes.