Fix status poller test

Author: toddmeng-db

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}");
         }