diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java
index d577b8614078..68dbd61f15f6 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java
@@ -784,7 +784,6 @@ public boolean advance() throws CoderException {
       flushStateInternal();
       Work newWork = additionalWork.work();
       ++workItemsPolled;
-      checkStateNotNull(keyTransitionListener).onKeyTransition(activeWork, newWork);
       startForNewKey(newWork);
       return true;
     }
diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java
index c5efcea4e47c..eb6bb51e4207 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java
@@ -550,11 +550,13 @@ public void testAdvance_success() throws Exception {
         .thenReturn(executableWork2)
         .thenReturn(null);
 
-    executionContext.start(
-        work1, workExecutor, mockExecutor, mockHandle, null, (oldWork, newWork) -> {});
+    StreamingModeExecutionContext.KeyTransitionListener mockListener =
+        mock(StreamingModeExecutionContext.KeyTransitionListener.class);
+    executionContext.start(work1, workExecutor, mockExecutor, mockHandle, null, mockListener);
 
     assertTrue(executionContext.advance());
     assertEquals("key2", executionContext.getSerializedKey().toStringUtf8());
+    verify(mockListener, times(1)).onKeyTransition(work1, work2);
     assertFalse(executionContext.advance());
   }