Changes to unit tests as per review