Fixing mock tests