From d4245e5caa5e9b2b6a4d53c331da4e12f26c10ff Mon Sep 17 00:00:00 2001
From: Madalyn Redding <66138537+m-redding@users.noreply.github.com>
Date: Thu, 17 Aug 2023 12:35:15 -0700
Subject: [PATCH] [Event Hubs] Partition Context fix (#38265)

* fix

* test fix
---
 .../EventHubListener.PartitionProcessor.cs       |  1 +
 .../tests/EventHubListenerTests.cs               | 16 ++++++++--------
 2 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/sdk/eventhub/Microsoft.Azure.WebJobs.Extensions.EventHubs/src/Listeners/EventHubListener.PartitionProcessor.cs b/sdk/eventhub/Microsoft.Azure.WebJobs.Extensions.EventHubs/src/Listeners/EventHubListener.PartitionProcessor.cs
index 89fb922c578f0..756774dd5cf0b 100644
--- a/sdk/eventhub/Microsoft.Azure.WebJobs.Extensions.EventHubs/src/Listeners/EventHubListener.PartitionProcessor.cs
+++ b/sdk/eventhub/Microsoft.Azure.WebJobs.Extensions.EventHubs/src/Listeners/EventHubListener.PartitionProcessor.cs
@@ -282,6 +282,7 @@ private async Task MonitorCachedEvents(DateTimeOffset? lastCheckpointTime, Cance
                         var details = GetOperationDetails(_mostRecentPartitionContext, "MaxWaitTimeElapsed");
                         _logger.LogDebug($"Partition Processor has waited MaxWaitTime since last invocation and is attempting to invoke function on all held events ({details})");
 
+                        UpdateCheckpointContext(triggerEvents, _mostRecentPartitionContext);
                         await TriggerExecute(triggerEvents, _mostRecentPartitionContext, _functionExecutionToken).ConfigureAwait(false);
                         if (!backgroundCancellationTokenSource.Token.IsCancellationRequested)
                         {
diff --git a/sdk/eventhub/Microsoft.Azure.WebJobs.Extensions.EventHubs/tests/EventHubListenerTests.cs b/sdk/eventhub/Microsoft.Azure.WebJobs.Extensions.EventHubs/tests/EventHubListenerTests.cs
index be11955f9b0e9..c42a28cdd22d7 100644
--- a/sdk/eventhub/Microsoft.Azure.WebJobs.Extensions.EventHubs/tests/EventHubListenerTests.cs
+++ b/sdk/eventhub/Microsoft.Azure.WebJobs.Extensions.EventHubs/tests/EventHubListenerTests.cs
@@ -112,10 +112,10 @@ public async Task ProcessEvents_MultipleDispatch_CheckpointsCorrectly(int batchC
                 Times.Exactly(expected));
         }
 
-        [TestCase(1, 29)]
-        [TestCase(5, 5)]
-        [TestCase(10, 2)]
-        [TestCase(30, 0)]
+        [TestCase(1, 30)]
+        [TestCase(5, 6)]
+        [TestCase(10, 3)]
+        [TestCase(30, 1)]
         [TestCase(35, 0)]
         public async Task ProcessEvents_MultipleDispatch_MinBatch_CheckpointsCorrectly_NoCheckpoint(int batchCheckpointFrequency, int expected)
         {
@@ -221,10 +221,10 @@ public async Task ProcessEvents_MultipleDispatch_MinBatch_CheckpointsCorrectly_R
                 Times.AtLeastOnce);
         }
 
-        [TestCase(1, 29)]
-        [TestCase(5, 5)]
-        [TestCase(10, 2)]
-        [TestCase(30, 0)]
+        [TestCase(1, 30)]
+        [TestCase(5, 6)]
+        [TestCase(10, 3)]
+        [TestCase(30, 1)]
         [TestCase(35, 0)]
         public async Task ProcessEvents_MultipleDispatch_MinBatch_CheckpointsCorrectly_OldCheckpoint(int batchCheckpointFrequency, int expected)
         {