Add unit tests for pkg/sharding/consistenthash (#480)

* Rename existing benchmark test file

* Add more comments

* Add unit tests for `pkg/sharding/consistenthash`

Author: timebertt

pkg/sharding/consistenthash/benchmark_test.go

@@ -0,0 +1,76 @@
+/*
+Copyright 2023 Tim Ebert.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package consistenthash
+
+import (
+	"fmt"
+	"math"
+	"testing"
+)
+
+func TestDistribution(t *testing.T) {
+	ring := New(DefaultHash, DefaultTokensPerNode)
+
+	hosts := generateHostnames(10)
+	dist := make(map[string]float64, len(hosts))
+	ring.AddNodes(hosts...)
+	for _, host := range hosts {
+		dist[host] = 0
+	}
+
+	// fmt.Println("Virtual Nodes:")
+	last := ring.tokens[len(ring.tokens)-1]
+	for _, token := range ring.tokens {
+		node := ring.tokenToNode[token]
+		percentage := float64(token-last) / math.MaxUint64
+		dist[node] += percentage
+
+		// fmt.Printf("\t%016x (%.5f): %.5f -> %s\n", token, float64(token)/math.MaxUint64, percentage, node)
+		last = token
+	}
+
+	fmt.Println("Nodes distribution:")
+	for _, host := range hosts {
+		fmt.Printf("\t%s: %.5f\n", host, dist[host])
+	}
+}
+
+func generateHostnames(n int) []string {
+	hosts := make([]string, n)
+	for i := range hosts {
+		host := fmt.Sprintf("10.42.0.%d", i)
+		hosts[i] = host
+	}
+	return hosts
+}
+
+func benchmarkRing(nodes int, tokensPerNode int, b *testing.B) {
+	hosts := generateHostnames(nodes)
+	b.ResetTimer()
+
+	for n := 0; n < b.N; n++ {
+		ring := New(DefaultHash, tokensPerNode, hosts...)
+		ring.Hash("Website.webhosting.timebertt.dev/project-foo/homepage")
+	}
+}
+
+func BenchmarkRing3_100(b *testing.B)   { benchmarkRing(3, 100, b) }
+func BenchmarkRing3_1000(b *testing.B)  { benchmarkRing(3, 1000, b) }
+func BenchmarkRing5_100(b *testing.B)   { benchmarkRing(5, 100, b) }
+func BenchmarkRing5_1000(b *testing.B)  { benchmarkRing(5, 1000, b) }
+func BenchmarkRing10_100(b *testing.B)  { benchmarkRing(10, 100, b) }
+func BenchmarkRing10_1000(b *testing.B) { benchmarkRing(10, 1000, b) }

pkg/sharding/consistenthash/consistenthash_suite_test.go

@@ -0,0 +1,29 @@
+/*
+Copyright 2025 Tim Ebert.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package consistenthash_test
+
+import (
+	"testing"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+func TestConsistentHash(t *testing.T) {
+	RegisterFailHandler(Fail)
+	RunSpecs(t, "Consistent Hash Suite")
+}

pkg/sharding/consistenthash/ring.go

@@ -32,18 +32,21 @@ var DefaultHash Hash = xxhash.Sum64String
 // DefaultTokensPerNode is the default number of virtual nodes per node.
 const DefaultTokensPerNode = 100
 
-// New creates a new hash ring.
-func New(fn Hash, tokensPerNode int, initialNodes ...string) *Ring {
-	if fn == nil {
-		fn = DefaultHash
+// New creates a new hash ring with the given configuration and adds the given nodes.
+// The given Hash (or DefaultHash if nil) is used to hash nodes and keys (strings).
+// Each node is assigned tokensPerNode tokens (or DefaultTokensPerNode if <= 0) – aka. virtual nodes – for a more
+// uniform key distribution.
+func New(hash Hash, tokensPerNode int, initialNodes ...string) *Ring {
+	if hash == nil {
+		hash = DefaultHash
 	}
 	if tokensPerNode <= 0 {
 		tokensPerNode = DefaultTokensPerNode
 	}
 
 	numTokens := len(initialNodes) * tokensPerNode
 	r := &Ring{
-		hash:          fn,
+		hash:          hash,
 		tokensPerNode: tokensPerNode,
 
 		tokens:      make([]uint64, 0, numTokens),
@@ -53,8 +56,9 @@ func New(fn Hash, tokensPerNode int, initialNodes ...string) *Ring {
 	return r
 }
 
-// Ring implements consistent hashing, aka ring hash (not thread-safe).
-// It hashes nodes and keys onto a ring of tokens. Keys are mapped to the next node on the ring.
+// Ring implements consistent hashing, aka. ring hash (not thread-safe).
+// It hashes nodes and keys (strings) onto a ring of tokens. Keys are mapped to the next token (node) on the ring.
+// Nodes cannot be removed. Instantiate a new Ring instead.
 type Ring struct {
 	hash          Hash
 	tokensPerNode int
@@ -63,10 +67,12 @@ type Ring struct {
 	tokenToNode map[uint64]string
 }
 
+// IsEmpty returns true if there are no nodes in this Ring.
 func (r *Ring) IsEmpty() bool {
 	return len(r.tokens) == 0
 }
 
+// AddNodes adds hash tokens for the given nodes to this Ring.
 func (r *Ring) AddNodes(nodes ...string) {
 	for _, node := range nodes {
 		for i := 0; i < r.tokensPerNode; i++ {
@@ -80,6 +86,7 @@ func (r *Ring) AddNodes(nodes ...string) {
 	slices.Sort(r.tokens)
 }
 
+// Hash hashes the given key onto the ring of tokens and returns the node that belongs to the next token on the ring.
 func (r *Ring) Hash(key string) string {
 	if r.IsEmpty() {
 		return ""

pkg/sharding/consistenthash/ring_test.go

@@ -1,5 +1,5 @@
 /*
-Copyright 2023 Tim Ebert.
+Copyright 2025 Tim Ebert.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -14,63 +14,78 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package consistenthash
+package consistenthash_test
 
 import (
-	"fmt"
-	"math"
-	"testing"
+	"strings"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+
+	. "github.com/timebertt/kubernetes-controller-sharding/pkg/sharding/consistenthash"
 )
 
-func TestDistribution(t *testing.T) {
-	ring := New(DefaultHash, DefaultTokensPerNode)
-
-	hosts := generateHostnames(10)
-	dist := make(map[string]float64, len(hosts))
-	ring.AddNodes(hosts...)
-	for _, host := range hosts {
-		dist[host] = 0
-	}
-
-	// fmt.Println("Virtual Nodes:")
-	last := ring.tokens[len(ring.tokens)-1]
-	for _, token := range ring.tokens {
-		node := ring.tokenToNode[token]
-		percentage := float64(token-last) / math.MaxUint64
-		dist[node] += percentage
-
-		// fmt.Printf("\t%016x (%.5f): %.5f -> %s\n", token, float64(token)/math.MaxUint64, percentage, node)
-		last = token
-	}
-
-	fmt.Println("Nodes distribution:")
-	for _, host := range hosts {
-		fmt.Printf("\t%s: %.5f\n", host, dist[host])
-	}
-}
-
-func generateHostnames(n int) []string {
-	hosts := make([]string, n)
-	for i := range hosts {
-		host := fmt.Sprintf("10.42.0.%d", i)
-		hosts[i] = host
-	}
-	return hosts
-}
-
-func benchmarkRing(nodes int, tokensPerNode int, b *testing.B) {
-	hosts := generateHostnames(nodes)
-	b.ResetTimer()
-
-	for n := 0; n < b.N; n++ {
-		ring := New(DefaultHash, tokensPerNode, hosts...)
-		ring.Hash("Website.webhosting.timebertt.dev/project-foo/homepage")
-	}
-}
-
-func BenchmarkRing3_100(b *testing.B)   { benchmarkRing(3, 100, b) }
-func BenchmarkRing3_1000(b *testing.B)  { benchmarkRing(3, 1000, b) }
-func BenchmarkRing5_100(b *testing.B)   { benchmarkRing(5, 100, b) }
-func BenchmarkRing5_1000(b *testing.B)  { benchmarkRing(5, 1000, b) }
-func BenchmarkRing10_100(b *testing.B)  { benchmarkRing(10, 100, b) }
-func BenchmarkRing10_1000(b *testing.B) { benchmarkRing(10, 1000, b) }
+var _ = Describe("Ring", func() {
+	Describe("#New", func() {
+		It("should initialize a new Ring", func() {
+			ring := New(nil, 0, "foo")
+			Expect(ring).NotTo(BeNil())
+			Expect(ring.IsEmpty()).To(BeFalse())
+		})
+	})
+
+	Describe("#IsEmpty", func() {
+		It("should true if there are no nodes", func() {
+			ring := New(nil, 0)
+			Expect(ring.IsEmpty()).To(BeTrue())
+			ring.AddNodes("foo")
+			Expect(ring.IsEmpty()).To(BeFalse())
+		})
+	})
+
+	Describe("#Hash", func() {
+		It("should use the configured hash function", func() {
+			ring := New(func(data string) uint64 {
+				if strings.HasPrefix(data, "foo") {
+					// map all foo* nodes and keys to 1
+					return 1
+				}
+				return 2
+			}, 1, "foo", "bar")
+
+			Expect(ring.Hash("foo")).To(Equal("foo"))
+			Expect(ring.Hash("bar")).To(Equal("bar"))
+			Expect(ring.Hash("baz")).To(Equal("bar"))
+		})
+
+		It("should use the default hash function", func() {
+			ring := New(nil, 0, "foo", "bar")
+
+			Expect(ring.Hash("1")).NotTo(Equal(ring.Hash("10")))
+		})
+
+		It("should return the empty string if there are no nodes", func() {
+			ring := New(nil, 0)
+
+			Expect(ring.Hash("foo")).To(BeEmpty())
+		})
+
+		It("should return the first node when walking the whole ring", func() {
+			ring := New(func(data string) uint64 {
+				if strings.HasPrefix(data, "foo") {
+					// map all foo* nodes and keys to 1
+					return 1
+				}
+				if strings.HasPrefix(data, "bar") {
+					// map all bar* nodes and keys to 1
+					return 2
+				}
+				return 3
+			}, 1, "foo", "bar")
+
+			Expect(ring.Hash("foo")).To(Equal("foo"))
+			Expect(ring.Hash("bar")).To(Equal("bar"))
+			Expect(ring.Hash("baz")).To(Equal("foo"))
+		})
+	})
+})