Merge pull request Automattic#15099 from Automattic/vkarpov15/Automatticgh-15088

fix(discriminator): gather childSchemas when creating discriminator to ensure $getAllSubdocs() can properly get all subdocs

Author: vkarpov15

lib/helpers/model/discriminator.js

@@ -119,6 +119,7 @@ module.exports = function discriminator(model, name, schema, tiedValue, applyPlu
     // schema. `Schema.prototype.clone()` copies `obj` by reference, no cloning.
     schema.obj = { ...schema.obj };
     mergeDiscriminatorSchema(schema, baseSchema);
+    schema._gatherChildSchemas();
 
     // Clean up conflicting paths _after_ merging re: gh-6076
     for (const conflictingPath of conflictingPaths) {

lib/schema.js

@@ -441,7 +441,7 @@ Schema.prototype._clone = function _clone(Constructor) {
       }
     }
   }
-  s.childSchemas = gatherChildSchemas(s);
+  s._gatherChildSchemas();
 
   s.virtuals = clone(this.virtuals);
   s.$globalPluginsApplied = this.$globalPluginsApplied;
@@ -1238,20 +1238,21 @@ Schema.prototype.path = function(path, obj) {
  * ignore
  */
 
-function gatherChildSchemas(schema) {
+Schema.prototype._gatherChildSchemas = function _gatherChildSchemas() {
   const childSchemas = [];
 
-  for (const path of Object.keys(schema.paths)) {
-    const schematype = schema.paths[path];
+  for (const path of Object.keys(this.paths)) {
+    const schematype = this.paths[path];
     if (schematype.$isMongooseDocumentArray || schematype.$isSingleNested) {
       childSchemas.push({ schema: schematype.schema, model: schematype.caster, path: path });
     } else if (schematype.$isSchemaMap && schematype.$__schemaType.$isSingleNested) {
       childSchemas.push({ schema: schematype.$__schemaType.schema, model: schematype.$__schemaType.caster, path: path });
     }
   }
 
+  this.childSchemas = childSchemas;
   return childSchemas;
-}
+};
 
 /*!
  * ignore

test/model.discriminator.test.js

@@ -2251,4 +2251,70 @@ describe('model', function() {
     assert(result);
 
   });
+
+  it('correctly gathers subdocs with discriminators (gh-15088)', async function() {
+    const RequestTriggeredWorkflowSchema = new Schema(
+      {
+        workflow: {
+          type: String
+        }
+      },
+      {
+        timestamps: true
+      }
+    );
+
+    const RequestSchema = new Schema(
+      {
+        status: {
+          type: String
+        },
+        summary: {
+          type: String
+        },
+        triggeredWorkflows: [RequestTriggeredWorkflowSchema]
+      },
+      {
+        discriminatorKey: 'source',
+        timestamps: true
+      }
+    );
+
+    const EmailRequestSchema = new Schema({
+      // Any extra properties that specifically apply to 'email' requests
+    });
+    const FormRequestSchema = new Schema({
+      // Any extra properties that specifically apply to 'form' requests
+    });
+
+    const Request = db.model('Request', RequestSchema);
+
+    Request.discriminator('Request:email', EmailRequestSchema, 'email');
+    Request.discriminator('Request:form', FormRequestSchema, 'form');
+
+    const request = await Request.create({
+      status: 'new',
+      source: 'form'
+    });
+
+    let requestAsReadFromDb = await Request.findById(request._id);
+
+    request.status = 'in progress';
+    await request.save();
+
+    assert.equal(requestAsReadFromDb.$getAllSubdocs().length, 0);
+    const now = new Date();
+    request.triggeredWorkflows.push({
+      workflow: '111111111111111111111111'
+    });
+    assert.equal(request.$getAllSubdocs().length, 1);
+    assert.equal(request.$getAllSubdocs()[0], request.triggeredWorkflows[0]);
+    await request.save();
+
+    requestAsReadFromDb = await Request.findById(request._id);
+    assert.equal(requestAsReadFromDb.$getAllSubdocs().length, 1);
+    assert.equal(requestAsReadFromDb.$getAllSubdocs()[0], requestAsReadFromDb.triggeredWorkflows[0]);
+    assert.ok(requestAsReadFromDb.triggeredWorkflows[0].createdAt.valueOf() >= now.valueOf());
+    assert.ok(requestAsReadFromDb.triggeredWorkflows[0].updatedAt.valueOf() >= now.valueOf());
+  });
 });