Refactor entity cleanup test to make it less flakey (#2871)

test/Common/DurableTaskEndToEndTests.cs

@@ -5058,6 +5058,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);
@@ -5077,28 +5079,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]