diff --git a/README.md b/README.md
index 9d95463..5e0cdf2 100644
--- a/README.md
+++ b/README.md
@@ -132,7 +132,7 @@ See [HttpMock](#httpmock) and [GraphQlMock](#graphqlmock) for more details.
 
 ### HttpResponseFunction
 
-> `function({ query, body, params, context, updateContext, getContext }): response | Promise<response>`
+> `function({ query, body, params, context, updateContext }): response | Promise<response>`
 
 <!-- https://www.tablesgenerator.com/markdown_tables -->
 
@@ -142,8 +142,7 @@ See [HttpMock](#httpmock) and [GraphQlMock](#graphqlmock) for more details.
 | body | `object` | `{}` | body object as defined by `express`. |
 | params | `object` | `{}` | params object as defined by `express`. |
 | context | `object` | `{}` | Data stored across API calls. |
-| updateContext | `Function` | `partialContext => updatedContext` | Used to update context. |
-| getContext | `Function` | `() => context` | Used to get the latest context. |
+| updateContext | `Function` | `partialContext => updatedContext` | Used to update context. `partialContext` can either be an `object` or a function (`context` => `partialContext`).  |
 | response | `undefined` / `Response` / `Override` | _required_ | [Response](#response), [Override](#override). |
 
 ### GraphQlMock
@@ -179,7 +178,7 @@ See [HttpMock](#httpmock) and [GraphQlMock](#graphqlmock) for more details.
 
 ### GraphQlResponseFunction
 
-> `function({ variables, context, updateContext, getContext }): response | Promise<response>`
+> `function({ variables, context, updateContext }): response | Promise<response>`
 
 <!-- https://www.tablesgenerator.com/markdown_tables -->
 
@@ -187,8 +186,7 @@ See [HttpMock](#httpmock) and [GraphQlMock](#graphqlmock) for more details.
 |----------|------|---------|-------------|
 | variables | `object` | `{}` | variables sent by client. |
 | context | `object` | `{}` | Data stored across API calls. |
-| updateContext | `Function` | `partialContext => updatedContext` | Used to update context. |
-| getContext | `Function` | `() => context` | Used to get the latest context. |
+| updateContext | `Function` | `partialContext => updatedContext` | Used to update context. `partialContext` can either be an `object` or a function (`context` => `partialContext`).  |
 | response | `undefined` / `Response` / `GraphQlResponse` / `Override` | _required_ | [Response](#response), [GraphQlResponse](#graphqlresponse), [Override](#override). |
 
 ### Override
diff --git a/src/index.spec.ts b/src/index.spec.ts
index d3c1dfa..f4eb4c1 100644
--- a/src/index.spec.ts
+++ b/src/index.spec.ts
@@ -892,6 +892,61 @@ describe('run', () => {
       });
     });
 
+    it('partial context can be set using a function', async () => {
+      const name = 'Betty';
+      const initialAge = 40;
+      const intervalDelayMs = 200;
+      const intervalTickCount = 5;
+      const timeoutDelayMs = intervalDelayMs * intervalTickCount + 100;
+
+      const server = run({
+        default: {
+          context: { age: initialAge, name },
+          mocks: [
+            {
+              url: '/info',
+              method: 'GET',
+              response: ({ context }) => context,
+            },
+            {
+              url: '/user',
+              method: 'POST',
+              response: ({ updateContext }) => {
+                const interval = setInterval(() => {
+                  updateContext(({ age }: any) => ({ age: age + 1 }));
+                }, intervalDelayMs);
+                setTimeout(() => {
+                  clearInterval(interval);
+                }, timeoutDelayMs);
+
+                return null;
+              },
+            },
+          ],
+        },
+      });
+
+      await serverTest(server, async () => {
+        const info1 = await rp.get('http://localhost:3000/info', {
+          json: true,
+        });
+        expect(info1).toEqual({ name, age: initialAge });
+
+        await rp.post('http://localhost:3000/user');
+
+        await new Promise(resolve => {
+          setTimeout(() => {
+            resolve();
+          }, timeoutDelayMs + 100);
+        });
+
+        const info2 = await rp.get('http://localhost:3000/info', {
+          json: true,
+        });
+        expect(info2).toEqual({ name, age: initialAge + intervalTickCount });
+      });
+    });
+
     it('context works for GraphQL requests', async () => {
       const initialName = 'Alice';
       const updatedName = 'Bob';
diff --git a/src/index.ts b/src/index.ts
index 6c2931e..671863e 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -126,8 +126,15 @@ function createRouter({
 
   return router;
 
-  function updateContext(partialContext: Context) {
-    context = { ...context, ...partialContext };
+  function updateContext(
+    partialContext: Context | ((context: Context) => Context),
+  ) {
+    context = {
+      ...context,
+      ...(typeof partialContext === 'function'
+        ? partialContext(context)
+        : partialContext),
+    };
 
     return context;
   }
diff --git a/src/types.ts b/src/types.ts
index b1b4cbc..b19243d 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -89,4 +89,6 @@ export type Options = {
 
 export type Context = Record<string, any>;
 
-export type UpdateContext = (partialContext: Context) => Context;
+export type UpdateContext = (
+  partialContext: Context | ((context: Context) => Context),
+) => Context;