Merge pull request #14343 from Automattic/vkarpov15/gh-14249

fix(populate): handle ref() functions that return a model instance
Automattic · Feb 12, 2024 · 37dc2f0 · 37dc2f0
2 parents c9877b6 + 3d0e643
commit 37dc2f0
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 1 deletion.
diff --git a/lib/helpers/populate/getModelsMapForPopulate.js b/lib/helpers/populate/getModelsMapForPopulate.js
@@ -494,7 +494,7 @@ function addModelNamesToMap(model, map, available, modelNames, options, data, re
 
   let k = modelNames.length;
   while (k--) {
-    const modelName = modelNames[k];
+    let modelName = modelNames[k];
     if (modelName == null) {
       continue;
     }
@@ -504,6 +504,7 @@ function addModelNamesToMap(model, map, available, modelNames, options, data, re
       Model = options.model;
     } else if (modelName[modelSymbol]) {
       Model = modelName;
+      modelName = Model.modelName;
     } else {
       try {
         Model = _getModelFromConn(connection, modelName);

diff --git a/test/model.populate.test.js b/test/model.populate.test.js
@@ -10922,4 +10922,50 @@ describe('model: populate:', function() {
     assert.equal(personFromDb.userType, 'User');
     assert.equal(personFromDb.userId.toHexString(), user._id.toHexString());
   });
+
+  it('handles ref() function that returns a model (gh-14249)', async function() {
+    const aSchema = new Schema({
+      name: String
+    });
+
+    const bSchema = new Schema({
+      name: String
+    });
+
+    const CategoryAModel = db.model('Test', aSchema);
+    const CategoryBModel = db.model('Test1', bSchema);
+
+    const testSchema = new Schema({
+      category: String,
+      subdoc: {
+        type: Schema.Types.ObjectId,
+        ref: function() {
+          return this.category === 'catA' ? CategoryAModel : CategoryBModel;
+        }
+      }
+    });
+
+    const parentSchema = new Schema({
+      name: String,
+      children: [testSchema]
+    });
+    const Parent = db.model('Parent', parentSchema);
+
+    const A = await CategoryAModel.create({
+      name: 'A'
+    });
+    const B = await CategoryBModel.create({
+      name: 'B'
+    });
+
+    const doc = await Parent.create({
+      name: 'Parent',
+      children: [{ category: 'catA', subdoc: A._id }, { category: 'catB', subdoc: B._id }]
+    });
+
+    const res = await Parent.findById(doc._id).populate('children.subdoc');
+    assert.equal(res.children.length, 2);
+    assert.equal(res.children[0].subdoc.name, 'A');
+    assert.equal(res.children[1].subdoc.name, 'B');
+  });
 });