added unit test

Author: farvez10

kinesumer.go

@@ -693,16 +693,21 @@ func (k *Kinesumer) consumeOnce(stream string, shard *Shard) ([]*kinesis.Record,
 	}
 	defer k.nextIters[stream].Store(shard.ID, output.NextShardIterator) // Update iter.
 
-	// set the shard closed state.
-	// If shard has no data, NextShardIterator will be nil.
-	// Reference: https://docs.aws.amazon.com/cli/latest/reference/kinesis/get-records.html#output
-	shard.Closed = output.NextShardIterator == nil
+	shard.Closed = getShardStatus(output)
 
 	// outer function has the for loop that takes care of the empty records case
 	// so not needed to check it here.
 	return output.Records, shard.Closed
 }
 
+// getShardStatus returns whether the shard is closed or not.
+func getShardStatus(output *kinesis.GetRecordsOutput) bool {
+	// set the shard closed state.
+	// If shard has no data, NextShardIterator will be nil.
+	// Reference: https://docs.aws.amazon.com/cli/latest/reference/kinesis/get-records.html#output
+	return output.NextShardIterator == nil
+}
+
 func (k *Kinesumer) getNextShardIterator(
 	ctx context.Context, stream, shardID string,
 ) (*string, error) {

kinesumer_test.go

@@ -682,3 +682,43 @@ func TestKinesumer_cleanupOffsetsWorksFine(t *testing.T) {
 		})
 	}
 }
+
+func Test_getShardStatus(t *testing.T) {
+	tests := []struct {
+		name  string // description of this test case
+		input struct {
+			output *kinesis.GetRecordsOutput
+		}
+		want bool
+	}{
+		{
+			name: "when shard is closed",
+			input: struct {
+				output *kinesis.GetRecordsOutput
+			}{
+				output: &kinesis.GetRecordsOutput{
+					NextShardIterator: nil,
+				},
+			},
+			want: true,
+		},
+
+		{
+			name: "when shard is open",
+			input: struct {
+				output *kinesis.GetRecordsOutput
+			}{
+				output: &kinesis.GetRecordsOutput{
+					NextShardIterator: aws.String("shardIterator"),
+				},
+			},
+			want: false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := getShardStatus(tt.input.output)
+			assert.Equal(t, tt.want, got, "Should be equal")
+		})
+	}
+}