From f3e47a04cef08173be32ae563a3fb7bbfe108936 Mon Sep 17 00:00:00 2001 From: Vitaly Terentyev Date: Wed, 16 Oct 2024 13:22:09 +0400 Subject: [PATCH 1/2] Fix assert for at least once method --- .../org/apache/beam/it/gcp/bigquery/BigQueryIOST.java | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) 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 22ff94e293b6..f05218cced43 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 int STORAGE_API_AT_LEAST_ONCE_MAX_ALLOWED_DIFFERENCE = 10_000; private static BigQueryResourceManager resourceManager; private static String tableName; @@ -336,9 +337,12 @@ private void generateDataAndWrite(BigQueryIO.Write writeIO) throws IOExc if (configuration.writeMethod.equals(STORAGE_API_AT_LEAST_ONCE_METHOD)) { 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", + (long) numRecords - rowCount, + STORAGE_API_AT_LEAST_ONCE_MAX_ALLOWED_DIFFERENCE, + rowCount, + (long) numRecords), + (long) numRecords - rowCount <= STORAGE_API_AT_LEAST_ONCE_MAX_ALLOWED_DIFFERENCE); } else { assertTrue( String.format( From 325092a3efbbf1edcfad4a85a21dc1eac7e5c854 Mon Sep 17 00:00:00 2001 From: Vitaly Terentyev Date: Wed, 16 Oct 2024 19:28:54 +0400 Subject: [PATCH 2/2] Refactoring --- .../apache/beam/it/gcp/bigquery/BigQueryIOST.java | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) 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 f05218cced43..ddb300d74f66 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,7 +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 int STORAGE_API_AT_LEAST_ONCE_MAX_ALLOWED_DIFFERENCE = 10_000; + private static final double STORAGE_API_AT_LEAST_ONCE_MAX_ALLOWED_DIFFERENCE_FRACTION = 0.00001; private static BigQueryResourceManager resourceManager; private static String tableName; @@ -335,14 +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( "Row difference (%d) exceeds the limit of %d. Rows: %d, Expected: %d", - (long) numRecords - rowCount, - STORAGE_API_AT_LEAST_ONCE_MAX_ALLOWED_DIFFERENCE, - rowCount, - (long) numRecords), - (long) numRecords - rowCount <= STORAGE_API_AT_LEAST_ONCE_MAX_ALLOWED_DIFFERENCE); + actualDifference, allowedDifference, rowCount, (long) numRecords), + actualDifference <= allowedDifference); } else { assertTrue( String.format(