From e496895d5e7fb5db51ae00f30f8ad965759c9f34 Mon Sep 17 00:00:00 2001 From: toddmeng-db Date: Wed, 27 Aug 2025 16:02:44 -0700 Subject: [PATCH] Fix status poller test --- csharp/test/Drivers/Databricks/E2E/StatementTests.cs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/csharp/test/Drivers/Databricks/E2E/StatementTests.cs b/csharp/test/Drivers/Databricks/E2E/StatementTests.cs index 07520b3a61..fe4f1711f2 100644 --- a/csharp/test/Drivers/Databricks/E2E/StatementTests.cs +++ b/csharp/test/Drivers/Databricks/E2E/StatementTests.cs @@ -908,9 +908,9 @@ private async Task PrepareTableAsync(string fullTableName, string sqlResourceLoc // NOTE: this is a thirty minute test. As of writing, databricks commands have 20 minutes of idle time (and checked every 5 minutes) [SkippableTheory] - [InlineData(false, "CloudFetch disabled")] - [InlineData(true, "CloudFetch enabled")] - public async Task StatusPollerKeepsQueryAlive(bool useCloudFetch, string configName) + [InlineData(false, "CloudFetch disabled", 30000000)] + [InlineData(true, "CloudFetch enabled", 3000000000)] + public async Task StatusPollerKeepsQueryAlive(bool useCloudFetch, string configName, long rowCount) { Skip.If(TestConfiguration.IsCITesting, "Skip test in CI testing"); @@ -925,7 +925,7 @@ public async Task StatusPollerKeepsQueryAlive(bool useCloudFetch, string configN using var statement = connection.CreateStatement(); // Execute a query that should return data - using a larger dataset to ensure multiple batches - statement.SqlQuery = "SELECT id, CAST(id AS STRING) as id_string, id * 2 as id_doubled FROM RANGE(30000000)"; + statement.SqlQuery = $"SELECT id, CAST(id AS STRING) as id_string, id * 2 as id_doubled FROM RANGE({rowCount})"; QueryResult result = statement.ExecuteQuery(); Assert.NotNull(result.Stream); @@ -950,7 +950,7 @@ public async Task StatusPollerKeepsQueryAlive(bool useCloudFetch, string configN } // Verify we got all rows - Assert.Equal(30000000, totalRows); + Assert.Equal(rowCount, totalRows); Assert.True(batchCount > 1, "Should have read multiple batches"); OutputHelper?.WriteLine($"Successfully read {totalRows} rows in {batchCount} batches after 30 minute delay with {configName}"); }