Improve test

apache · Feb 4, 2025 · 55f2d45 · 55f2d45
1 parent 373a716
commit 55f2d45
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/.../src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkStateTest.java b/.../src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkStateTest.java
@@ -473,11 +473,15 @@ public void testFailWork() {
           for (long cacheToken : Arrays.asList(5L, 6L)) {
             WorkId workId =
                 WorkId.builder().setWorkToken(++workToken).setCacheToken(cacheToken).build();
-            activeWorkState.failWorkForKey(
-                ImmutableList.of(WorkIdWithShardingKey.create(shardingKey, workId)));
+            boolean shouldFail = ThreadLocalRandom.current().nextBoolean();
+            if (shouldFail) {
+              activeWorkState.failWorkForKey(
+                  ImmutableList.of(WorkIdWithShardingKey.create(shardingKey, workId)));
+            }
             LinkedHashMap<WorkId, ExecutableWork> workIdExecutableWorkLinkedHashMap =
                 readOnlyActiveWork.get(shardingKey);
-            assertTrue(workIdExecutableWorkLinkedHashMap.get(workId).work().isFailed());
+            assertEquals(
+                shouldFail, workIdExecutableWorkLinkedHashMap.get(workId).work().isFailed());
           }
         }
       }