add CHash.RemoveAllGroup()

werbenhu · Apr 14, 2023 · 888dde9 · 888dde9
1 parent 63aeee4
commit 888dde9
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 0 deletions.
diff --git a/chash.go b/chash.go
@@ -48,6 +48,15 @@ func (c *CHash) RemoveGroup(groupName string) {
 	delete(c.groups, groupName)
 }
 
+func (c *CHash) RemoveAllGroup() {
+	c.Lock()
+	defer c.Unlock()
+
+	for k := range c.groups {
+		delete(c.groups, k)
+	}
+}
+
 func (c *CHash) Insert(groupName string, key string, payload []byte) error {
 	c.Lock()
 	group, ok := c.groups[groupName]

diff --git a/chash_test.go b/chash_test.go
@@ -83,6 +83,21 @@ func TestCHashRemoveGroup(t *testing.T) {
 	assert.Equal(t, ErrGroupNotFound, err)
 }
 
+func TestCHashRemoveAllGroup(t *testing.T) {
+	hash := New()
+	group1, err := hash.CreateGroup("werbenhu1", 2000)
+	assert.Nil(t, err)
+	assert.NotNil(t, group1)
+
+	group2, err := hash.CreateGroup("werbenhu2", 1000)
+	assert.Nil(t, err)
+	assert.NotNil(t, group2)
+
+	assert.Equal(t, 2, len(hash.groups))
+	hash.RemoveAllGroup()
+	assert.Equal(t, 0, len(hash.groups))
+}
+
 func TestCHashInsert(t *testing.T) {
 	hash := New()
 	hash.CreateGroup("werbenhu1", 10000)