From 77f56794c9c6bda472ca9ce7787a2d1e10e97024 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Sun, 25 Jun 2023 20:33:30 -0700
Subject: [PATCH] types/key: add test for NodePublic.Shard

Updates #cleanup

Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 types/key/node_test.go | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/types/key/node_test.go b/types/key/node_test.go
index 26f580e97..edbd9e927 100644
--- a/types/key/node_test.go
+++ b/types/key/node_test.go
@@ -157,3 +157,26 @@ func TestChallenge(t *testing.T) {
 		t.Errorf("didn't round trip: %v != %v", back, pub)
 	}
 }
+
+// Test that NodePublic.Shard is uniformly distributed.
+func TestShard(t *testing.T) {
+	const N = 1_000
+	var shardCount [256]int
+	for i := 0; i < N; i++ {
+		shardCount[NewNode().Public().Shard()]++
+	}
+	e := float64(N) / 256 // expected
+	var x2 float64        // chi-squared
+	for _, c := range shardCount {
+		r := float64(c) - e // residual
+		x2 += r * r / e
+	}
+	t.Logf("x^2 = %v", x2)
+	if x2 > 512 { // really want x^2 =~ (256 - 1), but leave slop
+		t.Errorf("too much variation in shard distribution")
+		for i, c := range shardCount {
+			rj := float64(c) - e
+			t.Logf("shard[%v] = %v (off by %v)", i, c, rj)
+		}
+	}
+}