From 33cb4f4e66ee46dd9eede856c73d7d8e2f1f465d Mon Sep 17 00:00:00 2001
From: Zijun Wang <zijunw@amazon.com>
Date: Thu, 2 Nov 2023 14:53:07 -0700
Subject: [PATCH] Address PR comments

---
 pkg/utils/common.go      |  3 ++
 pkg/utils/common_test.go | 90 +++++++++++++++++++++++++++++-----------
 2 files changed, 69 insertions(+), 24 deletions(-)

diff --git a/pkg/utils/common.go b/pkg/utils/common.go
index ef8a51e9..72e87000 100644
--- a/pkg/utils/common.go
+++ b/pkg/utils/common.go
@@ -20,6 +20,9 @@ func min[T constraints.Ordered](a, b T) T {
 
 func Chunks[T any](in []T, size int) [][]T {
 	out := [][]T{}
+	if size <= 0 {
+		return out
+	}
 	for low := 0; low < len(in); low += size {
 		high := min(low+size, len(in))
 		out = append(out, in[low:high])
diff --git a/pkg/utils/common_test.go b/pkg/utils/common_test.go
index 081001d6..32c115a4 100644
--- a/pkg/utils/common_test.go
+++ b/pkg/utils/common_test.go
@@ -1,36 +1,78 @@
 package utils
 
 import (
+	"reflect"
 	"testing"
-
-	"github.com/stretchr/testify/assert"
 )
 
 func TestChunks(t *testing.T) {
-	size := 2
 	obj1 := &struct{}{}
 	obj2 := &struct{}{}
 	obj3 := &struct{}{}
+	obj4 := &struct{}{}
+	obj5 := &struct{}{}
 
-	t.Run("[obj1] -> [[obj1]]", func(t *testing.T) {
-
-		in := []*struct{}{obj1}
-		out := Chunks(in, size)
-		assert.Equal(t, [][]*struct{}{{obj1}}, out)
-
-	})
-
-	t.Run("[obj1,obj2] -> [[obj1,obj2]]", func(t *testing.T) {
-
-		in := []*struct{}{obj1, obj2}
-		got := Chunks(in, size)
-		assert.Equal(t, [][]*struct{}{{obj1, obj2}}, got)
-
-	})
-
-	t.Run("[obj1, obj2, obj3] -> [[obj1,obj2],[obj3]]", func(t *testing.T) {
-		in := []*struct{}{obj1, obj2, obj3}
-		got := Chunks(in, size)
-		assert.Equal(t, [][]*struct{}{{obj1, obj2}, {obj3}}, got)
-	})
+	type testCase struct {
+		name string
+		in   []*struct{}
+		size int
+		want [][]*struct{}
+	}
+	tests := []testCase{
+		{
+			name: "[] -> []",
+			in:   []*struct{}{},
+			size: 2,
+			want: [][]*struct{}{},
+		},
+		{
+			name: "[obj1] -> [[obj1]]",
+			in:   []*struct{}{obj1},
+			size: 2,
+			want: [][]*struct{}{{obj1}},
+		},
+		{
+			name: "[obj1, obj2] -> [[obj1,obj2]]",
+			in:   []*struct{}{obj1},
+			size: 2,
+			want: [][]*struct{}{{obj1}},
+		},
+		{
+			name: "[obj1, obj2, obj3] -> [[obj1,obj2],[obj3]]",
+			in:   []*struct{}{obj1, obj2, obj3},
+			size: 2,
+			want: [][]*struct{}{{obj1, obj2}, {obj3}},
+		},
+		{
+			name: "[obj1, obj2, obj3, obj4] -> [[obj1,obj2],[obj3,obj4]]",
+			in:   []*struct{}{obj1, obj2, obj3, obj4},
+			size: 2,
+			want: [][]*struct{}{{obj1, obj2}, {obj3, obj4}},
+		},
+		{
+			name: "[obj1, obj2, obj3, obj4, obj5] -> [[obj1,obj2],[obj3,obj4],[obj5]]",
+			in:   []*struct{}{obj1, obj2, obj3, obj4, obj5},
+			size: 2,
+			want: [][]*struct{}{{obj1, obj2}, {obj3, obj4}, {obj5}},
+		},
+		{
+			name: "size 0",
+			in:   []*struct{}{obj1, obj2, obj3, obj4, obj5},
+			size: 0,
+			want: [][]*struct{}{},
+		},
+		{
+			name: "negative size",
+			in:   []*struct{}{obj1, obj2, obj3, obj4, obj5},
+			size: -1,
+			want: [][]*struct{}{},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := Chunks(tt.in, tt.size); !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("Chunks() got %v, want %v", got, tt.want)
+			}
+		})
+	}
 }