Refactor entity cleanup test to make it less flakey

test/Common/DurableTaskEndToEndTests.cs

@@ -5065,6 +5065,8 @@ public async Task DurableEntity_CleanEntityStorage(string storageProvider)
                 var orchestrationA = $"{prefix}-A";
                 var orchestrationB = $"{prefix}-B";
 
+                // PART 1: Test removal of empty entities
+
                 // create an empty entity
                 var client = await host.StartOrchestratorAsync(nameof(TestOrchestrations.CreateEmptyEntities), new EntityId[] { emptyEntityId }, this.output);
                 var status = await client.WaitForCompletionAsync(this.output);
@@ -5084,28 +5086,38 @@ public async Task DurableEntity_CleanEntityStorage(string storageProvider)
                 var result = await client.InnerClient.ListEntitiesAsync(query, CancellationToken.None);
                 Assert.Contains(result.Entities, s => s.EntityId.Equals(emptyEntityId));
 
+                // test removal of empty entity
+                var response = await client.InnerClient.CleanEntityStorageAsync(removeEmptyEntities: true, releaseOrphanedLocks: false, CancellationToken.None);
+                Assert.Equal(1, response.NumberOfEmptyEntitiesRemoved);
+                Assert.Equal(0, response.NumberOfOrphanedLocksRemoved);
+
+                // check that the empty entity record has been removed from storage
+                result = await client.InnerClient.ListEntitiesAsync(query, CancellationToken.None);
+                Assert.DoesNotContain(result.Entities, s => s.EntityId.Equals(emptyEntityId));
+
+                // PART 2: Test recovery from orphaned locks
+
                 // run an orchestration A that leaves an orphaned lock
                 TestDurableClient clientA = await host.StartOrchestratorAsync(nameof(TestOrchestrations.LockThenFailReplay), (orphanedEntityId, true), this.output, orchestrationA);
                 status = await clientA.WaitForCompletionAsync(this.output);
 
                 // run an orchestration B that queues behind A for the lock (and thus gets stuck)
                 TestDurableClient clientB = await host.StartOrchestratorAsync(nameof(TestOrchestrations.LockThenFailReplay), (orphanedEntityId, false), this.output, orchestrationB);
 
-                // remove empty entity and release orphaned lock
-                var response = await client.InnerClient.CleanEntityStorageAsync(true, true, CancellationToken.None);
+                await Task.Delay(TimeSpan.FromMinutes(1)); // wait for a stable entity executionID, needed until https://github.com/Azure/durabletask/pull/1128 is merged
+
+                // remove release orphaned lock to unblock orchestration B
+                // Note: do NOT remove empty entities yet: we want to keep the empty entity so it can unblock orchestration B
+                response = await client.InnerClient.CleanEntityStorageAsync(removeEmptyEntities: false, releaseOrphanedLocks: true, CancellationToken.None);
                 Assert.Equal(1, response.NumberOfOrphanedLocksRemoved);
-                Assert.Equal(1, response.NumberOfEmptyEntitiesRemoved);
+                Assert.Equal(0, response.NumberOfEmptyEntitiesRemoved);
 
                 // wait for orchestration B to complete, now that the lock has been released
                 status = await clientB.WaitForCompletionAsync(this.output);
                 Assert.True(status.RuntimeStatus == OrchestrationRuntimeStatus.Completed);
 
-                // check that the empty entity record has been removed from storage
-                result = await client.InnerClient.ListEntitiesAsync(query, CancellationToken.None);
-                Assert.DoesNotContain(result.Entities, s => s.EntityId.Equals(emptyEntityId));
-
                 // clean again to remove the orphaned entity which is now empty also
-                response = await client.InnerClient.CleanEntityStorageAsync(true, true, CancellationToken.None);
+                response = await client.InnerClient.CleanEntityStorageAsync(removeEmptyEntities: true, releaseOrphanedLocks: true, CancellationToken.None);
                 Assert.Equal(0, response.NumberOfOrphanedLocksRemoved);
                 Assert.Equal(1, response.NumberOfEmptyEntitiesRemoved);
 

test/Common/HttpApiHandlerTests.cs

@@ -411,14 +411,14 @@ public async Task WaitForCompletionOrCreateCheckStatusResponseAsync_Returns_HTTP
                     TaskHub = TestConstants.TaskHub,
                     ConnectionName = TestConstants.ConnectionName,
                 },
-                TimeSpan.FromSeconds(10),
+                TimeSpan.FromSeconds(15),
                 TimeSpan.FromSeconds(3));
             stopwatch.Stop();
             Assert.Equal(HttpStatusCode.OK, httpResponseMessage.StatusCode);
             var content = await httpResponseMessage.Content.ReadAsStringAsync();
             var value = JsonConvert.DeserializeObject<string>(content);
             Assert.Equal("Hello Tokyo!", value);
-            Assert.True(stopwatch.Elapsed < TimeSpan.FromSeconds(10));
+            Assert.True(stopwatch.Elapsed < TimeSpan.FromSeconds(15));
         }
 
         [Fact]