Context not properly provided through data injector and subscriptions

packages/react-relay/__tests__/LiveResolvers-test.js

@@ -1876,3 +1876,88 @@ test('ResolverContext can contain observable values', async () => {
     counter_context: 1,
   });
 });
+
+test('ResolverContext as passed through nested resolver counters', async () => {
+  const source = RelayRecordSource.create({
+    'client:root': {
+      __id: 'client:root',
+      __typename: '__Root',
+      me: {__ref: '1'},
+    },
+    '1': {
+      __id: '1',
+      __typename: 'User',
+      id: '1',
+    },
+  });
+  const FooQuery = graphql`
+    query LiveResolversTestCounterContextBaseQuery {
+      base_counter_context {
+        count_plus_one
+      }
+    }
+  `;
+
+  let next: (v: number) => void = () => {
+    throw new Error('next() not initialized');
+  };
+
+  const operation = createOperationDescriptor(FooQuery, {});
+  const store = new RelayModernStore(source, {
+    gcReleaseBufferSize: 0,
+    resolverContext: {
+      counter: Observable.create<number>(observer => {
+        next = (value: number) => observer.next(value);
+      }),
+    },
+  });
+
+  const environment = new RelayModernEnvironment({
+    network: RelayNetwork.create(jest.fn()),
+    store,
+  });
+
+  let observedCounter = null;
+
+  const snapshot = environment.lookup(operation.fragment);
+  // $FlowFixMe[unclear-type] - lookup() doesn't have the nice types of reading a fragment through the actual APIs:
+  observedCounter = (snapshot.data: any).base_counter_context.count_plus_one;
+
+  const environmentUpdateHandler = jest.fn(() => {
+    const s = environment.lookup(operation.fragment);
+    // $FlowFixMe[unclear-type] - lookup() doesn't have the nice types of reading a fragment through the actual APIs:
+    observedCounter = (s.data: any).base_counter_context.count_plus_one;
+  });
+  const disposable = environment.subscribe(
+    snapshot,
+    // $FlowFixMe[invalid-tuple-arity] Error found while enabling LTI on this file
+    environmentUpdateHandler,
+  );
+
+  // SETUP COMPLETE
+
+  // Read the initial value
+  expect(observedCounter).toBe(0);
+  expect(environmentUpdateHandler).not.toHaveBeenCalled();
+
+  // Increment and assert we get notified of the new value
+  next(43);
+  expect(environmentUpdateHandler).toHaveBeenCalledTimes(1);
+  expect(observedCounter).toBe(44);
+
+  // Unsubscribe then increment and assert don't get notified.
+  disposable.dispose();
+  next(1);
+  expect(environmentUpdateHandler).toHaveBeenCalledTimes(1);
+  expect(observedCounter).toBe(44);
+
+  // Explicitly read and assert we see the incremented value
+  // missed before due to unsubscribing.
+  const nextSnapshot = environment.lookup(operation.fragment);
+
+  expect(nextSnapshot.data).toEqual({
+    base_counter_context: {
+      count_plus_one: 2,
+    },
+  });
+});

packages/react-relay/__tests__/RelayResolverModelWithContext-test.js

@@ -0,0 +1,134 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @flow strict-local
+ * @format
+ * @oncall relay
+ */
+
+'use strict';
+
+import type {RelayResolverModelWithContextTestFragment$key} from './__generated__/RelayResolverModelWithContextTestFragment.graphql';
+import type {TestResolverContextType} from 'relay-runtime/mutations/__tests__/TestResolverContextType';
+
+const React = require('react');
+const {
+  RelayEnvironmentProvider,
+  useClientQuery,
+  useFragment,
+} = require('react-relay');
+const TestRenderer = require('react-test-renderer');
+const {Observable} = require('relay-runtime');
+const RelayNetwork = require('relay-runtime/network/RelayNetwork');
+const {graphql} = require('relay-runtime/query/GraphQLTag');
+const {
+  addTodo,
+  resetStore,
+} = require('relay-runtime/store/__tests__/resolvers/ExampleTodoStore');
+const RelayModernEnvironment = require('relay-runtime/store/RelayModernEnvironment');
+const RelayModernStore = require('relay-runtime/store/RelayModernStore.js');
+const RelayRecordSource = require('relay-runtime/store/RelayRecordSource');
+const {
+  disallowConsoleErrors,
+  disallowWarnings,
+  injectPromisePolyfill__DEPRECATED,
+} = require('relay-test-utils-internal');
+
+injectPromisePolyfill__DEPRECATED();
+disallowWarnings();
+disallowConsoleErrors();
+
+beforeEach(() => {
+  resetStore(() => {});
+});
+
+function EnvironmentWrapper({
+  children,
+  environment,
+}: {
+  children: React.Node,
+  environment: RelayModernEnvironment,
+}) {
+  return (
+    <RelayEnvironmentProvider environment={environment}>
+      <React.Suspense fallback="Loading...">{children}</React.Suspense>
+    </RelayEnvironmentProvider>
+  );
+}
+
+const RelayResolverModelWithContextTestFragment = graphql`
+  fragment RelayResolverModelWithContextTestFragment on TodoModel {
+    id
+    description
+    another_value_from_context
+  }
+`;
+
+const RelayResolverModelWithContextTestQuery = graphql`
+  query RelayResolverModelWithContextTestQuery($id: ID!) {
+    todo_model(todoID: $id) {
+      ...RelayResolverModelWithContextTestFragment
+    }
+  }
+`;
+
+describe('RelayResolverModelWithContext', () => {
+  function TodoComponent(props: {
+    fragmentKey: ?RelayResolverModelWithContextTestFragment$key,
+  }) {
+    const data = useFragment(
+      RelayResolverModelWithContextTestFragment,
+      props.fragmentKey,
+    );
+    if (data == null) {
+      return null;
+    }
+
+    return data.another_value_from_context;
+  }
+
+  function TodoRootComponent(props: {todoID: string}) {
+    const data = useClientQuery(RelayResolverModelWithContextTestQuery, {
+      id: props.todoID,
+    });
+    if (data?.todo_model == null) {
+      return null;
+    }
+
+    return <TodoComponent fragmentKey={data?.todo_model} />;
+  }
+
+  test('returns a value from context when resolverDataInjector is used', () => {
+    const resolverContext: TestResolverContextType = {
+      greeting: {
+        myHello: 'This is a value from context',
+      },
+      counter: Observable.create<number>(observer => {
+        observer.next(10);
+      }),
+    };
+
+    const store = new RelayModernStore(RelayRecordSource.create(), {
+      resolverContext,
+    });
+    const environment = new RelayModernEnvironment({
+      network: RelayNetwork.create(jest.fn()),
+      store,
+    });
+
+    addTodo('Test todo');
+
+    let renderer;
+    TestRenderer.act(() => {
+      renderer = TestRenderer.create(
+        <EnvironmentWrapper environment={environment}>
+          <TodoRootComponent todoID="todo-1" />
+        </EnvironmentWrapper>,
+      );
+    });
+    expect(renderer?.toJSON()).toEqual('This is a value from context');
+  });
+});