Ajay/Melanie | #2889/#2891 | Fixing tests