Add an in-memory checkpoint to the API

consumer_test.go

@@ -3,6 +3,7 @@ package consumer
 import (
 	"context"
 	"fmt"
+	"github.com/harlow/kinesis-consumer/store/memory"
 	"sync"
 	"testing"
 
@@ -52,7 +53,7 @@ func TestScan(t *testing.T) {
 		},
 	}
 	var (
-		cp  = &fakeCheckpoint{cache: map[string]string{}}
+		cp  = memory.New()
 		ctr = &fakeCounter{}
 	)
 
@@ -114,7 +115,7 @@ func TestScanShard(t *testing.T) {
 	}
 
 	var (
-		cp  = &fakeCheckpoint{cache: map[string]string{}}
+		cp  = memory.New()
 		ctr = &fakeCounter{}
 	)
 
@@ -219,7 +220,7 @@ func TestScanShard_SkipCheckpoint(t *testing.T) {
 		},
 	}
 
-	var cp = &fakeCheckpoint{cache: map[string]string{}}
+	var cp = memory.New()
 
 	c, err := New("myStreamName", WithClient(client), WithStore(cp))
 	if err != nil {
@@ -335,29 +336,6 @@ func (c *kinesisClientMock) GetShardIteratorWithContext(ctx aws.Context, in *kin
 	return c.getShardIteratorMock(in)
 }
 
-// implementation of checkpoint
-type fakeCheckpoint struct {
-	cache map[string]string
-	mu    sync.Mutex
-}
-
-func (fc *fakeCheckpoint) SetCheckpoint(streamName, shardID, sequenceNumber string) error {
-	fc.mu.Lock()
-	defer fc.mu.Unlock()
-
-	key := fmt.Sprintf("%s-%s", streamName, shardID)
-	fc.cache[key] = sequenceNumber
-	return nil
-}
-
-func (fc *fakeCheckpoint) GetCheckpoint(streamName, shardID string) (string, error) {
-	fc.mu.Lock()
-	defer fc.mu.Unlock()
-
-	key := fmt.Sprintf("%s-%s", streamName, shardID)
-	return fc.cache[key], nil
-}
-
 // implementation of counter
 type fakeCounter struct {
 	counter int64

store/memory/store.go

@@ -0,0 +1,33 @@
+// The memory store provides a store that can be used for testing and single-threaded applications.
+// DO NOT USE this in a production application where persistence beyond a single application lifecycle is necessary
+// or when there are multiple consumers.
+package memory
+
+import (
+  "fmt"
+  "sync"
+)
+
+func New() *Checkpoint{
+  return &Checkpoint{}
+}
+
+type Checkpoint struct {
+  sync.Map
+}
+
+func (c *Checkpoint) SetCheckpoint(streamName, shardID, sequenceNumber string) error {
+  if sequenceNumber == "" {
+    return fmt.Errorf("sequence number should not be empty")
+  }
+  c.Store(streamName+":"+shardID, sequenceNumber)
+  return nil
+}
+
+func (c *Checkpoint) GetCheckpoint(streamName, shardID string) (string, error) {
+  val, ok := c.Load(streamName + ":" + shardID)
+  if !ok {
+    return "", nil
+  }
+  return val.(string), nil
+}

store/memory/store_test.go

@@ -0,0 +1,32 @@
+package memory_test
+
+import (
+  "testing"
+
+  "github.com/harlow/kinesis-consumer/store/memory"
+)
+
+func Test_CheckpointLifecycle(t *testing.T) {
+  c := memory.New()
+
+  // set
+  c.SetCheckpoint("streamName", "shardID", "testSeqNum")
+
+  // get
+  val, err := c.GetCheckpoint("streamName", "shardID")
+  if err != nil  {
+    t.Fatalf("get checkpoint error: %v", err)
+  }
+  if val != "testSeqNum" {
+    t.Fatalf("checkpoint exists expected %s, got %s", "testSeqNum", val)
+  }
+}
+
+func Test_SetEmptySeqNum(t *testing.T) {
+  c := memory.New()
+
+  err := c.SetCheckpoint("streamName", "shardID", "")
+  if err == nil || err.Error() != "sequence number should not be empty" {
+    t.Fatalf("should not allow empty sequence number")
+  }
+}