create failing test

package.json

@@ -5,7 +5,7 @@
   "main": "index.js",
   "scripts": {
     "pretest": "npm run compile",
-    "test": "mocha --require ./test/fixtures/setup.js --reporter spec --full-trace --recursive ./lib/test && jest",
+    "test": "mocha --require ./test/fixtures/setup.js --reporter spec --full-trace --recursive ./lib/test/react-web/server && jest",
     "posttest": "npm run lint",
     "filesize": "npm run compile:browser && ./scripts/filesize.js --file=./dist/index.min.js --maxGzip=15",
     "compile": "tsc",

src/graphql.tsx

@@ -319,7 +319,7 @@ export default function graphql(
 
       setInitialProps() {
         if (this.type === DocumentType.Mutation) {
-           this.createWrappedMutation(this.props);
+          this.createWrappedMutation(this.props);
           return;
         }
 
@@ -346,7 +346,6 @@ export default function graphql(
               opts.query = document;
               return this.client.query(opts);
             };
-
             queryData = assign({
               errors: null, loading: false, variables, refetch, fetchMore,
             }, result);

src/server.ts

@@ -34,7 +34,6 @@ export function getChildFromComponent(component) {
 function getQueriesFromTree(
   { component, context = {}, queries = []}: QueryTreeArgument, fetch: boolean = true
 ) {
-
   if (!component) return;
 
   // stateless function
@@ -50,18 +49,17 @@ function getQueriesFromTree(
 
     let newContext = context;
     if (Component.getChildContext) newContext = assign({}, context, Component.getChildContext());
-
     context = newContext;
 
     // see if there is a fetch data method
     if (typeof type.fetchData === 'function' && fetch) {
-      const query = type.fetchData(ownProps, newContext);
+      const query = type.fetchData(ownProps, context);
       if (query) queries.push({ query, component });
     }
 
     getQueriesFromTree({
       component: getChildFromComponent(Component),
-      context: newContext,
+      context,
       queries,
     });
   } else if (props && props.children) {
@@ -71,14 +69,14 @@ function getQueriesFromTree(
       queries,
     }));
   }
-
   return { queries, context };
 }
 
 // XXX component Cache
 export function getDataFromTree(app, ctx: any = {}, fetch: boolean = true): Promise<any> {
 
   let { context, queries } = getQueriesFromTree({ component: app, context: ctx }, fetch);
+
   // no queries found, nothing to do
   if (!queries.length) return Promise.resolve(context);
 

test/react-web/server/index.tsx

@@ -60,6 +60,57 @@ describe('SSR', () => {
         ;
     });
 
+    it('should not require `ApolloProvider` to be the root component', (done) => {
+
+      const query = gql`{ currentUser { firstName } }`;
+      const data = { currentUser: { firstName: 'James' } };
+      const networkInterface = mockNetworkInterface(
+        { request: { query }, result: { data }, delay: 50 }
+      );
+      const apolloClient = new ApolloClient({ networkInterface });
+
+      const WrappedElement = graphql(query)(({ data }) => (
+        <div>{data.loading ? 'loading' : data.currentUser.firstName}</div>
+      ));
+
+      class MyRootContainer extends React.Component<any, any> {
+
+        constructor(props) {
+          super(props);
+          this.state = { color: 'purple' };
+        }
+
+        getChildContext() {
+          return { color: this.state.color };
+        }
+
+        render() {
+          return <div>{this.props.children}</div>;
+        }
+      }
+
+      (MyRootContainer as any).childContextTypes = {
+        color: React.PropTypes.string,
+      };
+
+      const app = (
+        <MyRootContainer>
+          <ApolloProvider client={apolloClient}>
+            <WrappedElement />
+          </ApolloProvider>
+        </MyRootContainer>
+      );
+
+      getDataFromTree(app)
+        .then(() => {
+          const markup = ReactDOM.renderToString(app);
+          expect(markup).to.match(/James/);
+          done();
+        })
+        .catch(done)
+        ;
+    });
+
     it('should run return the initial state for hydration', (done) => {
       const query = gql`{ currentUser { firstName } }`;
       const data = { currentUser: { firstName: 'James' } };