Skip to content

Commit bfaf7b9

Browse files
committed
Add additional window function tests for better coverage
Add 4 new test cases covering different window function scenarios: - DENSE_RANK() for dense ranking functionality - Multiple window functions in single query (ROW_NUMBER, RANK, COUNT) - COUNT(*) and AVG() OVER () for aggregate window functions - CTE combined with window functions These tests complement the existing window function tests and ensure comprehensive coverage of window function behavior.
1 parent 79df09d commit bfaf7b9

File tree

7 files changed

+179
-3
lines changed

7 files changed

+179
-3
lines changed

internal/stackql/driver/window_cte_integration_test.go

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -219,3 +219,139 @@ func TestSelectComputeDisksCTEMultiple(t *testing.T) {
219219
stackqltestutil.SetupSimpleSelectGoogleComputeDisks(t, 2)
220220
stackqltestutil.RunCaptureTestAgainstFiles(t, testSubject, []string{testobjects.ExpectedSelectComputeDisksCTEMultiple})
221221
}
222+
223+
//nolint:govet,lll // test file
224+
func TestSelectComputeDisksWindowDenseRank(t *testing.T) {
225+
runtimeCtx, err := stackqltestutil.GetRuntimeCtx(testobjects.GetGoogleProviderString(), "text", "TestSelectComputeDisksWindowDenseRank")
226+
if err != nil {
227+
t.Fatalf("Test failed: %v", err)
228+
}
229+
inputBundle, err := stackqltestutil.BuildInputBundle(*runtimeCtx)
230+
if err != nil {
231+
t.Fatalf("Test failed: %v", err)
232+
}
233+
234+
testSubject := func(t *testing.T, outFile *bufio.Writer) {
235+
handlerCtx, err := entryutil.BuildHandlerContext(*runtimeCtx, strings.NewReader(""), lrucache.NewLRUCache(int64(runtimeCtx.QueryCacheSize)), inputBundle.WithStdOut(outFile), true)
236+
if err != nil {
237+
t.Fatalf("Test failed: %v", err)
238+
}
239+
240+
handlerCtx.SetQuery(testobjects.SelectGoogleComputeDisksWindowDenseRank)
241+
dr, _ := NewStackQLDriver(handlerCtx)
242+
querySubmitter := querysubmit.NewQuerySubmitter()
243+
prepareErr := querySubmitter.PrepareQuery(handlerCtx)
244+
if prepareErr != nil {
245+
t.Fatalf("Test failed: %v", prepareErr)
246+
}
247+
response := querySubmitter.SubmitQuery()
248+
responsehandler.HandleResponse(handlerCtx, response)
249+
250+
dr.ProcessQuery(handlerCtx.GetRawQuery())
251+
}
252+
253+
stackqltestutil.SetupSimpleSelectGoogleComputeDisks(t, 1)
254+
stackqltestutil.RunCaptureTestAgainstFiles(t, testSubject, []string{testobjects.ExpectedSelectComputeDisksWindowDenseRank})
255+
}
256+
257+
//nolint:govet,lll // test file
258+
func TestSelectComputeDisksWindowMultiple(t *testing.T) {
259+
runtimeCtx, err := stackqltestutil.GetRuntimeCtx(testobjects.GetGoogleProviderString(), "text", "TestSelectComputeDisksWindowMultiple")
260+
if err != nil {
261+
t.Fatalf("Test failed: %v", err)
262+
}
263+
inputBundle, err := stackqltestutil.BuildInputBundle(*runtimeCtx)
264+
if err != nil {
265+
t.Fatalf("Test failed: %v", err)
266+
}
267+
268+
testSubject := func(t *testing.T, outFile *bufio.Writer) {
269+
handlerCtx, err := entryutil.BuildHandlerContext(*runtimeCtx, strings.NewReader(""), lrucache.NewLRUCache(int64(runtimeCtx.QueryCacheSize)), inputBundle.WithStdOut(outFile), true)
270+
if err != nil {
271+
t.Fatalf("Test failed: %v", err)
272+
}
273+
274+
handlerCtx.SetQuery(testobjects.SelectGoogleComputeDisksWindowMultiple)
275+
dr, _ := NewStackQLDriver(handlerCtx)
276+
querySubmitter := querysubmit.NewQuerySubmitter()
277+
prepareErr := querySubmitter.PrepareQuery(handlerCtx)
278+
if prepareErr != nil {
279+
t.Fatalf("Test failed: %v", prepareErr)
280+
}
281+
response := querySubmitter.SubmitQuery()
282+
responsehandler.HandleResponse(handlerCtx, response)
283+
284+
dr.ProcessQuery(handlerCtx.GetRawQuery())
285+
}
286+
287+
stackqltestutil.SetupSimpleSelectGoogleComputeDisks(t, 1)
288+
stackqltestutil.RunCaptureTestAgainstFiles(t, testSubject, []string{testobjects.ExpectedSelectComputeDisksWindowMultiple})
289+
}
290+
291+
//nolint:govet,lll // test file
292+
func TestSelectComputeDisksWindowCount(t *testing.T) {
293+
runtimeCtx, err := stackqltestutil.GetRuntimeCtx(testobjects.GetGoogleProviderString(), "text", "TestSelectComputeDisksWindowCount")
294+
if err != nil {
295+
t.Fatalf("Test failed: %v", err)
296+
}
297+
inputBundle, err := stackqltestutil.BuildInputBundle(*runtimeCtx)
298+
if err != nil {
299+
t.Fatalf("Test failed: %v", err)
300+
}
301+
302+
testSubject := func(t *testing.T, outFile *bufio.Writer) {
303+
handlerCtx, err := entryutil.BuildHandlerContext(*runtimeCtx, strings.NewReader(""), lrucache.NewLRUCache(int64(runtimeCtx.QueryCacheSize)), inputBundle.WithStdOut(outFile), true)
304+
if err != nil {
305+
t.Fatalf("Test failed: %v", err)
306+
}
307+
308+
handlerCtx.SetQuery(testobjects.SelectGoogleComputeDisksWindowCount)
309+
dr, _ := NewStackQLDriver(handlerCtx)
310+
querySubmitter := querysubmit.NewQuerySubmitter()
311+
prepareErr := querySubmitter.PrepareQuery(handlerCtx)
312+
if prepareErr != nil {
313+
t.Fatalf("Test failed: %v", prepareErr)
314+
}
315+
response := querySubmitter.SubmitQuery()
316+
responsehandler.HandleResponse(handlerCtx, response)
317+
318+
dr.ProcessQuery(handlerCtx.GetRawQuery())
319+
}
320+
321+
stackqltestutil.SetupSimpleSelectGoogleComputeDisks(t, 1)
322+
stackqltestutil.RunCaptureTestAgainstFiles(t, testSubject, []string{testobjects.ExpectedSelectComputeDisksWindowCount})
323+
}
324+
325+
//nolint:govet,lll // test file
326+
func TestSelectComputeDisksCTEWithWindow(t *testing.T) {
327+
runtimeCtx, err := stackqltestutil.GetRuntimeCtx(testobjects.GetGoogleProviderString(), "text", "TestSelectComputeDisksCTEWithWindow")
328+
if err != nil {
329+
t.Fatalf("Test failed: %v", err)
330+
}
331+
inputBundle, err := stackqltestutil.BuildInputBundle(*runtimeCtx)
332+
if err != nil {
333+
t.Fatalf("Test failed: %v", err)
334+
}
335+
336+
testSubject := func(t *testing.T, outFile *bufio.Writer) {
337+
handlerCtx, err := entryutil.BuildHandlerContext(*runtimeCtx, strings.NewReader(""), lrucache.NewLRUCache(int64(runtimeCtx.QueryCacheSize)), inputBundle.WithStdOut(outFile), true)
338+
if err != nil {
339+
t.Fatalf("Test failed: %v", err)
340+
}
341+
342+
handlerCtx.SetQuery(testobjects.SelectGoogleComputeDisksCTEWithWindow)
343+
dr, _ := NewStackQLDriver(handlerCtx)
344+
querySubmitter := querysubmit.NewQuerySubmitter()
345+
prepareErr := querySubmitter.PrepareQuery(handlerCtx)
346+
if prepareErr != nil {
347+
t.Fatalf("Test failed: %v", prepareErr)
348+
}
349+
response := querySubmitter.SubmitQuery()
350+
responsehandler.HandleResponse(handlerCtx, response)
351+
352+
dr.ProcessQuery(handlerCtx.GetRawQuery())
353+
}
354+
355+
stackqltestutil.SetupSimpleSelectGoogleComputeDisks(t, 1)
356+
stackqltestutil.RunCaptureTestAgainstFiles(t, testSubject, []string{testobjects.ExpectedSelectComputeDisksCTEWithWindow})
357+
}

internal/test/testobjects/expected.go

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -44,9 +44,13 @@ const (
4444
ExpectedSelectExecOrgGetIamPolicyAgg string = "test/assets/expected/aggregated-select/google/cloudresourcemanager/select-exec-getiampolicy-agg.csv"
4545

4646
// Window function expected outputs.
47-
ExpectedSelectComputeDisksWindowRowNumber string = "test/assets/expected/window-select/google/disks/text/disks-window-row-number.csv"
48-
ExpectedSelectComputeDisksWindowRank string = "test/assets/expected/window-select/google/disks/text/disks-window-rank.csv"
49-
ExpectedSelectComputeDisksWindowSum string = "test/assets/expected/window-select/google/disks/text/disks-window-sum.csv"
47+
ExpectedSelectComputeDisksWindowRowNumber string = "test/assets/expected/window-select/google/disks/text/disks-window-row-number.csv"
48+
ExpectedSelectComputeDisksWindowRank string = "test/assets/expected/window-select/google/disks/text/disks-window-rank.csv"
49+
ExpectedSelectComputeDisksWindowSum string = "test/assets/expected/window-select/google/disks/text/disks-window-sum.csv"
50+
ExpectedSelectComputeDisksWindowDenseRank string = "test/assets/expected/window-select/google/disks/text/disks-window-dense-rank.csv"
51+
ExpectedSelectComputeDisksWindowMultiple string = "test/assets/expected/window-select/google/disks/text/disks-window-multiple.csv"
52+
ExpectedSelectComputeDisksWindowCount string = "test/assets/expected/window-select/google/disks/text/disks-window-count.csv"
53+
ExpectedSelectComputeDisksCTEWithWindow string = "test/assets/expected/window-select/google/disks/text/disks-cte-with-window.csv"
5054

5155
// CTE expected outputs.
5256
ExpectedSelectComputeDisksCTESimple string = "test/assets/expected/cte-select/google/disks/text/disks-cte-simple.csv"

internal/test/testobjects/input.go

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,14 @@ const (
6161
SelectGoogleComputeDisksWindowRank string = `select name, sizeGb, RANK() OVER (ORDER BY sizeGb) as size_rank from google.compute.disks where zone = 'australia-southeast1-b' AND project = 'testing-project' ORDER BY name;`
6262
SelectGoogleComputeDisksWindowSum string = `select name, sizeGb, SUM(cast(sizeGb as unsigned)) OVER (ORDER BY name) as running_total from google.compute.disks where zone = 'australia-southeast1-b' AND project = 'testing-project' ORDER BY name;`
6363

64+
// Additional window function test queries.
65+
SelectGoogleComputeDisksWindowDenseRank string = `select name, sizeGb, DENSE_RANK() OVER (ORDER BY sizeGb) as dense_rank from google.compute.disks where zone = 'australia-southeast1-b' AND project = 'testing-project' ORDER BY name;`
66+
SelectGoogleComputeDisksWindowMultiple string = `select name, sizeGb, ROW_NUMBER() OVER (ORDER BY name) as row_num, RANK() OVER (ORDER BY sizeGb DESC) as size_rank, COUNT(*) OVER () as total_count from google.compute.disks where zone = 'australia-southeast1-b' AND project = 'testing-project' ORDER BY name;`
67+
SelectGoogleComputeDisksWindowCount string = `select name, sizeGb, COUNT(*) OVER () as total, AVG(cast(sizeGb as unsigned)) OVER () as avg_size from google.compute.disks where zone = 'australia-southeast1-b' AND project = 'testing-project' ORDER BY name;`
68+
69+
// CTE with window function test query.
70+
SelectGoogleComputeDisksCTEWithWindow string = `WITH disk_cte AS (SELECT name, sizeGb FROM google.compute.disks WHERE zone = 'australia-southeast1-b' AND project = 'testing-project') SELECT name, sizeGb, ROW_NUMBER() OVER (ORDER BY name) as row_num FROM disk_cte ORDER BY name;`
71+
6472
// CTE test queries.
6573
SelectGoogleComputeDisksCTESimple string = `WITH disk_cte AS (SELECT name, sizeGb FROM google.compute.disks WHERE zone = 'australia-southeast1-b' AND project = 'testing-project') SELECT name, sizeGb FROM disk_cte ORDER BY name;`
6674
SelectGoogleComputeDisksCTEWithAgg string = `WITH disk_cte AS (SELECT name, sizeGb FROM google.compute.disks WHERE zone = 'australia-southeast1-b' AND project = 'testing-project') SELECT COUNT(*) as disk_count FROM disk_cte;`
Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
name,sizeGb,row_num
2+
demo-disk-qq1,10,1
3+
demo-disk-qq2,10,2
4+
demo-disk-xx2,10,3
5+
demo-disk-xx3,20,4
6+
demo-disk-xx4,30,5
7+
demo-disk-xx5,40,6
Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
name,sizeGb,total,avg_size
2+
demo-disk-qq1,10,6,20.0
3+
demo-disk-qq2,10,6,20.0
4+
demo-disk-xx2,10,6,20.0
5+
demo-disk-xx3,20,6,20.0
6+
demo-disk-xx4,30,6,20.0
7+
demo-disk-xx5,40,6,20.0
Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
name,sizeGb,dense_rank
2+
demo-disk-qq1,10,1
3+
demo-disk-qq2,10,1
4+
demo-disk-xx2,10,1
5+
demo-disk-xx3,20,2
6+
demo-disk-xx4,30,3
7+
demo-disk-xx5,40,4
Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
name,sizeGb,row_num,size_rank,total_count
2+
demo-disk-qq1,10,1,4,6
3+
demo-disk-qq2,10,2,4,6
4+
demo-disk-xx2,10,3,4,6
5+
demo-disk-xx3,20,4,3,6
6+
demo-disk-xx4,30,5,2,6
7+
demo-disk-xx5,40,6,1,6

0 commit comments

Comments
 (0)