Fix getShardIds

consumer.go

@@ -241,21 +241,37 @@ func (c *Consumer) handleRecord(shardID string, r *Record, fn func(*Record) Scan
 }
 
 func (c *Consumer) getShardIDs(streamName string) ([]string, error) {
+	var ss []string
+	var exclusiveStartShardId *string
+	for {
 		resp, err := c.client.DescribeStream(
 			&kinesis.DescribeStreamInput{
 				StreamName:            aws.String(streamName),
+				ExclusiveStartShardId: exclusiveStartShardId,
 			},
 		)
 		if err != nil {
 			return nil, fmt.Errorf("describe stream error: %v", err)
 		}
 
-	var ss []string
+		streamDescription := resp.StreamDescription
-	for _, shard := range resp.StreamDescription.Shards {
+		shards := streamDescription.Shards
+
+		if len(shards) == 0 {
+			return ss, nil
+		}
+
+		for _, shard := range shards {
 			ss = append(ss, *shard.ShardId)
 		}
+
+		exclusiveStartShardId = shards[len(shards)-1].ShardId
+
+		if *streamDescription.HasMoreShards == false {
 			return ss, nil
 		}
+	}
+}
 
 func (c *Consumer) getShardIterator(streamName, shardID, lastSeqNum string) (*string, error) {
 	params := &kinesis.GetShardIteratorInput{

consumer_test.go

@@ -46,7 +46,7 @@ func TestConsumer_Scan(t *testing.T) {
 					Shards: []*kinesis.Shard{
 						{ShardId: aws.String("myShard")},
 					},
-				},
+					HasMoreShards: aws.Bool(false)},
 			}, nil
 		},
 	}