diff --git a/wayang-tests-integration/src/test/java/org/apache/wayang/tests/TensorflowIrisScalaLikeApiIT.java b/wayang-tests-integration/src/test/java/org/apache/wayang/tests/TensorflowIrisScalaLikeApiIT.java
index 39abf922a..47e23ec4f 100644
--- a/wayang-tests-integration/src/test/java/org/apache/wayang/tests/TensorflowIrisScalaLikeApiIT.java
+++ b/wayang-tests-integration/src/test/java/org/apache/wayang/tests/TensorflowIrisScalaLikeApiIT.java
@@ -160,8 +160,13 @@ void test() {
         DataQuantaBuilder<?, String> textFileSource = plan.readTextFile(uri.toString());
 
         if (random) {
+            // Assign random keys for shuffling, then sort by those keys
+            // This avoids comparator contract violations from r.nextInt()
             Random r = new Random();
-            textFileSource = textFileSource.sort(e -> r.nextInt());
+            textFileSource = textFileSource
+                    .map(line -> new Tuple<>(r.nextDouble(), line))
+                    .sort(tuple -> tuple.field0)
+                    .map(tuple -> tuple.field1);
         }
 
         MapDataQuantaBuilder<String, Tuple<float[], Integer>> mapXY = textFileSource.map(line -> {