Correctly report error for invalid root type in extended schema

src/type/__tests__/validation-test.js

@@ -323,6 +323,52 @@ describe('Type System: A Schema must have Object root types', () => {
     ]);
   });
 
+  it('rejects a schema extended with invalid root types', () => {
+    let schema = buildSchema(`
+      input SomeInputObject {
+        test: String
+      }
+    `);
+
+    schema = extendSchema(schema, parse(`
+      extend schema {
+        query: SomeInputObject
+      }
+    `));
+
+    schema = extendSchema(schema, parse(`
+      extend schema {
+        mutation: SomeInputObject
+      }
+    `));
+
+    schema = extendSchema(schema, parse(`
+      extend schema {
+        subscription: SomeInputObject
+      }
+    `));
+
+    const errorMsg = (operation) =>
+      `${operation} root type must be Object type, it cannot be SomeInputObject.`;
+    expect(validateSchema(schema)).to.deep.equal([
+      {
+        message:
+          'Query root type must be Object type, it cannot be SomeInputObject.',
+        locations: [{ line: 3, column: 16 }],
+      },
+      {
+        message:
+          'Mutation root type must be Object type if provided, it cannot be SomeInputObject.',
+        locations: [{ line: 3, column: 19 }],
+      },
+      {
+        message:
+          'Subscription root type must be Object type if provided, it cannot be SomeInputObject.',
+        locations: [{ line: 3, column: 23 }],
+      },
+    ]);
+  });
+
   it('rejects a Schema whose directives are incorrectly typed', () => {
     const schema = new GraphQLSchema({
       query: SomeObjectType,

src/type/validate.js

@@ -152,13 +152,15 @@ function getOperationTypeNode(
   type: GraphQLObjectType,
   operation: string,
 ): ?ASTNode {
-  const astNode = schema.astNode;
-  const operationTypeNode =
-    astNode &&
-    astNode.operationTypes.find(
-      operationType => operationType.operation === operation,
-    );
-  return operationTypeNode ? operationTypeNode.type : type && type.astNode;
+  for (const node of getAllNodes(schema)) {
+    for (const operationType of node.operationTypes) {
+      if (operationType.operation === operation) {
+        return operationType.type;
+      }
+    }
+  }
+
+  return type.astNode;
 }
 
 function validateDirectives(context: SchemaValidationContext): void {