diff --git a/x-pack/plugin/data-frame/src/test/java/org/elasticsearch/xpack/dataframe/checkpoint/DataFrameTransformsCheckpointServiceTests.java b/x-pack/plugin/data-frame/src/test/java/org/elasticsearch/xpack/dataframe/checkpoint/DataFrameTransformsCheckpointServiceTests.java
index 6c487765f92..ea5362d184b 100644
--- a/x-pack/plugin/data-frame/src/test/java/org/elasticsearch/xpack/dataframe/checkpoint/DataFrameTransformsCheckpointServiceTests.java
+++ b/x-pack/plugin/data-frame/src/test/java/org/elasticsearch/xpack/dataframe/checkpoint/DataFrameTransformsCheckpointServiceTests.java
@@ -82,7 +82,6 @@ public class DataFrameTransformsCheckpointServiceTests extends ESTestCase {
         }
     }
 
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/40368")
     public void testExtractIndexCheckpointsInconsistentGlobalCheckpoints() {
         Map<String, long[]> expectedCheckpoints = new HashMap<>();
         Set<String> indices = randomUserIndices();
@@ -141,9 +140,19 @@ public class DataFrameTransformsCheckpointServiceTests extends ESTestCase {
                 // we need at least one replica for testing
                 int numShardCopies = randomIntBetween(2, 4);
 
+                int primaryShard = 0;
+                if (skipPrimaries) {
+                    primaryShard = randomInt(numShardCopies - 1);
+                }
                 int inconsistentReplica = -1;
                 if (inconsistentGlobalCheckpoints) {
-                    inconsistentReplica = randomIntBetween(0, numShardCopies - 1);
+                    List<Integer> replicas = new ArrayList<>(numShardCopies - 1);
+                    for (int i = 0; i < numShardCopies; i++) {
+                        if (primaryShard != i) {
+                            replicas.add(i);
+                        }
+                    }
+                    inconsistentReplica = randomFrom(replicas);
                 }
 
                 // SeqNoStats asserts that checkpoints are logical
@@ -156,11 +165,7 @@ public class DataFrameTransformsCheckpointServiceTests extends ESTestCase {
 
                 for (int replica = 0;  replica < numShardCopies; replica++) {
                     ShardId shardId = new ShardId(index, shardIndex);
-                    boolean primary = (replica == 0);
-
-                    if (skipPrimaries) {
-                        primary = randomBoolean();
-                    }
+                    boolean primary = (replica == primaryShard);
 
                     Path path = createTempDir().resolve("indices").resolve(index.getUUID()).resolve(String.valueOf(shardIndex));
                     ShardRouting shardRouting = ShardRouting.newUnassigned(shardId, primary,