Merge pull request #14314 from Automattic/vkarpov15/gh-14285

fix(populate): call setter on virtual populated path with populated doc instead of undefined

lib/helpers/populate/setPopulatedVirtualValue.js

@@ -0,0 +1,33 @@
+'use strict';
+
+/**
+ * Set a populated virtual value on a document's `$$populatedVirtuals` value
+ *
+ * @param {*} populatedVirtuals A document's `$$populatedVirtuals`
+ * @param {*} name The virtual name
+ * @param {*} v The result of the populate query
+ * @param {*} options The populate options. This function handles `justOne` and `count` options.
+ * @returns {Array<Document>|Document|Object|Array<Object>} the populated virtual value that was set
+ */
+
+module.exports = function setPopulatedVirtualValue(populatedVirtuals, name, v, options) {
+  if (options.justOne || options.count) {
+    populatedVirtuals[name] = Array.isArray(v) ?
+      v[0] :
+      v;
+
+    if (typeof populatedVirtuals[name] !== 'object') {
+      populatedVirtuals[name] = options.count ? v : null;
+    }
+  } else {
+    populatedVirtuals[name] = Array.isArray(v) ?
+      v :
+      v == null ? [] : [v];
+
+    populatedVirtuals[name] = populatedVirtuals[name].filter(function(doc) {
+      return doc && typeof doc === 'object';
+    });
+  }
+
+  return populatedVirtuals[name];
+};

lib/schema.js

@@ -20,6 +20,7 @@ const handleReadPreferenceAliases = require('./helpers/query/handleReadPreferenc
 const idGetter = require('./helpers/schema/idGetter');
 const merge = require('./helpers/schema/merge');
 const mpath = require('mpath');
+const setPopulatedVirtualValue = require('./helpers/populate/setPopulatedVirtualValue');
 const setupTimestamps = require('./helpers/timestamps/setupTimestamps');
 const utils = require('./utils');
 const validateRef = require('./helpers/populate/validateRef');
@@ -2288,28 +2289,17 @@ Schema.prototype.virtual = function(name, options) {
     virtual.options = options;
 
     virtual.
-      set(function(_v) {
+      set(function(v) {
         if (!this.$$populatedVirtuals) {
           this.$$populatedVirtuals = {};
         }
 
-        if (options.justOne || options.count) {
-          this.$$populatedVirtuals[name] = Array.isArray(_v) ?
-            _v[0] :
-            _v;
-
-          if (typeof this.$$populatedVirtuals[name] !== 'object') {
-            this.$$populatedVirtuals[name] = options.count ? _v : null;
-          }
-        } else {
-          this.$$populatedVirtuals[name] = Array.isArray(_v) ?
-            _v :
-            _v == null ? [] : [_v];
-
-          this.$$populatedVirtuals[name] = this.$$populatedVirtuals[name].filter(function(doc) {
-            return doc && typeof doc === 'object';
-          });
-        }
+        return setPopulatedVirtualValue(
+          this.$$populatedVirtuals,
+          name,
+          v,
+          options
+        );
       });
 
     if (typeof options.get === 'function') {

test/model.populate.test.js

@@ -10868,4 +10868,58 @@ describe('model: populate:', function() {
       { name: 'foo', prop: 'bar' }
     );
   });
+
+  it('calls setter on virtual populated path with populated doc (gh-14285)', async function() {
+    const userSchema = new Schema({
+      email: String,
+      name: 'String'
+    });
+
+    const User = db.model('User', userSchema);
+
+    const user = await User.create({
+      email: 'admin@example.com',
+      name: 'Admin'
+    });
+
+    const personSchema = new Schema({
+      userId: ObjectId,
+      userType: String
+    });
+
+    personSchema.
+      virtual('user', {
+        ref() {
+          return this.userType;
+        },
+        localField: 'userId',
+        foreignField: '_id',
+        justOne: true
+      }).
+      set(function(user) {
+        if (user) {
+          this.userId = user._id;
+          this.userType = user.constructor.modelName;
+        } else {
+          this.userId = null;
+          this.userType = null;
+        }
+
+        return user;
+      });
+
+    const Person = db.model('Person', personSchema);
+
+    const person = new Person({
+      userId: user._id,
+      userType: 'User'
+    });
+
+    await person.save();
+
+    const personFromDb = await Person.findById(person._id).populate('user');
+    assert.equal(personFromDb.user.name, 'Admin');
+    assert.equal(personFromDb.userType, 'User');
+    assert.equal(personFromDb.userId.toHexString(), user._id.toHexString());
+  });
 });