Fix issues discovered while testing (OMS-396)