diff --git a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOST.java b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOST.java index 22ff94e293b6a..ddb300d74f667 100644 --- a/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOST.java +++ b/it/google-cloud-platform/src/test/java/org/apache/beam/it/gcp/bigquery/BigQueryIOST.java @@ -84,6 +84,7 @@ public final class BigQueryIOST extends IOStressTestBase { private static final String READ_ELEMENT_METRIC_NAME = "read_count"; private static final String STORAGE_WRITE_API_METHOD = "STORAGE_WRITE_API"; private static final String STORAGE_API_AT_LEAST_ONCE_METHOD = "STORAGE_API_AT_LEAST_ONCE"; + private static final double STORAGE_API_AT_LEAST_ONCE_MAX_ALLOWED_DIFFERENCE_FRACTION = 0.00001; private static BigQueryResourceManager resourceManager; private static String tableName; @@ -334,11 +335,14 @@ private void generateDataAndWrite(BigQueryIO.Write writeIO) throws IOExc // Depending on writing method there might be duplicates on different sides (read or write). if (configuration.writeMethod.equals(STORAGE_API_AT_LEAST_ONCE_METHOD)) { + long allowedDifference = + (long) (numRecords * STORAGE_API_AT_LEAST_ONCE_MAX_ALLOWED_DIFFERENCE_FRACTION); + long actualDifference = (long) numRecords - rowCount; assertTrue( String.format( - "Number of rows in the table (%d) is less than the expected number (%d). Missing records: %d", - rowCount, (long) numRecords, (long) numRecords - rowCount), - rowCount >= numRecords); + "Row difference (%d) exceeds the limit of %d. Rows: %d, Expected: %d", + actualDifference, allowedDifference, rowCount, (long) numRecords), + actualDifference <= allowedDifference); } else { assertTrue( String.format(