Add Size method in MapLimitStore

coinpaprika · Mar 5, 2019 · d3c6aae · d3c6aae
1 parent 4a7a0ce
commit d3c6aae
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 0 deletions.
diff --git a/map_limit_store.go b/map_limit_store.go
@@ -60,6 +60,13 @@ func (m *MapLimitStore) Get(key string, previousWindow, currentWindow time.Time)
 	return prevValue, currValue, nil
 }
 
+// Size returns current length of data map
+func (m *MapLimitStore) Size() int {
+	m.mutex.RLock()
+	defer m.mutex.RUnlock()
+	return len(m.data)
+}
+
 func mapKey(key string, window time.Time) string {
 	return fmt.Sprintf("%s_%s", key, window.Format(time.RFC3339))
 }
diff --git a/map_limit_store_test.go b/map_limit_store_test.go
@@ -92,3 +92,33 @@ func TestMapLimitStore_Get(t *testing.T) {
 		assert.Equal(t, tt.wantCurrValue, currVal)
 	}
 }
+
+func TestMapLimitStore_Size(t *testing.T) {
+	tests := []struct {
+		name   string
+		key    string
+		window time.Time
+		size   int
+	}{
+		{
+			name:   "test_MapLimitStore_Size",
+			key:    "tt",
+			window: time.Now().UTC(),
+			size:   1,
+		},
+		{
+			name:   "test_MapLimitStore_Size",
+			key:    "tt",
+			window: time.Time{},
+			size:   0,
+		},
+	}
+	for _, tt := range tests {
+		m := NewMapLimitStore(1*time.Minute, 10*time.Second)
+		if !tt.window.IsZero() {
+			err := m.Inc(tt.key, tt.window)
+			assert.NoError(t, err)
+		}
+		assert.Equal(t, tt.size, m.Size())
+	}
+}