From b73954cd7c54c3ef937646eacfbc316db26f5cfc Mon Sep 17 00:00:00 2001
From: Renan DelValle <commit@ridv.xyz>
Date: Sat, 14 Sep 2019 10:20:23 -0700
Subject: [PATCH] Adding batch calculator utility and test cases for it.

---
 util.go      | 19 +++++++++++++++++++
 util_test.go | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+)

diff --git a/util.go b/util.go
index 08f3490..989f8e8 100644
--- a/util.go
+++ b/util.go
@@ -99,3 +99,22 @@ func validateAuroraURL(location string) (string, error) {
 
 	return u.String(), nil
 }
+
+func calculateCurrentBatch(updatingInstances int32, batchSizes []int32) int {
+	for i, size := range batchSizes {
+		updatingInstances -= size
+		if updatingInstances <= 0 {
+			return i
+		}
+	}
+
+	// Overflow batches
+	batchCount := len(batchSizes) - 1
+	lastBatchIndex := len(batchSizes) - 1
+	batchCount += int(updatingInstances / batchSizes[lastBatchIndex])
+
+	if updatingInstances%batchSizes[lastBatchIndex] != 0 {
+		batchCount++
+	}
+	return batchCount
+}
diff --git a/util_test.go b/util_test.go
index e015535..b8341b2 100644
--- a/util_test.go
+++ b/util_test.go
@@ -63,3 +63,40 @@ func TestAuroraURLValidator(t *testing.T) {
 		assert.NoError(t, err)
 	})
 }
+
+func TestCurrentBatchCalculator(t *testing.T) {
+	t.Run("singleBatchOverflow", func(t *testing.T) {
+		curBatch := calculateCurrentBatch(10, []int32{2})
+		assert.Equal(t, 4, curBatch)
+	})
+
+	t.Run("noInstancesUpdating", func(t *testing.T) {
+		curBatch := calculateCurrentBatch(0, []int32{2})
+		assert.Equal(t, 0, curBatch)
+	})
+
+	t.Run("evenMatchSingleBatch", func(t *testing.T) {
+		curBatch := calculateCurrentBatch(2, []int32{2})
+		assert.Equal(t, 0, curBatch)
+	})
+
+	t.Run("moreInstancesThanBatches", func(t *testing.T) {
+		curBatch := calculateCurrentBatch(5, []int32{1, 2})
+		assert.Equal(t, 2, curBatch)
+	})
+
+	t.Run("moreInstancesThanBatchesDecreasing", func(t *testing.T) {
+		curBatch := calculateCurrentBatch(5, []int32{2, 1})
+		assert.Equal(t, 3, curBatch)
+	})
+
+	t.Run("unevenFit", func(t *testing.T) {
+		curBatch := calculateCurrentBatch(2, []int32{1, 2})
+		assert.Equal(t, 1, curBatch)
+	})
+
+	t.Run("halfWay", func(t *testing.T) {
+		curBatch := calculateCurrentBatch(1, []int32{1, 2})
+		assert.Equal(t, 0, curBatch)
+	})
+}