From f9c445bada8e58a5d63426d2ff3ac13ad3f5260b Mon Sep 17 00:00:00 2001
From: Ivan Kripakov <ivan.kripakov@pashapay.az>
Date: Thu, 4 Jan 2024 18:25:24 +0400
Subject: [PATCH] add simple server_lookup_test.go & small fix (use ID instead
 of name in server_lookup.go)

---
 dkron/server_lookup.go      |  2 +-
 dkron/server_lookup_test.go | 73 +++++++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+), 1 deletion(-)
 create mode 100644 dkron/server_lookup_test.go

diff --git a/dkron/server_lookup.go b/dkron/server_lookup.go
index eac4038c9..2343cc433 100644
--- a/dkron/server_lookup.go
+++ b/dkron/server_lookup.go
@@ -26,7 +26,7 @@ func (sl *ServerLookup) AddServer(server *ServerParts) {
 	sl.lock.Lock()
 	defer sl.lock.Unlock()
 	sl.addressToServer[raft.ServerAddress(server.RPCAddr.String())] = server
-	sl.idToServer[raft.ServerID(server.Name)] = server
+	sl.idToServer[raft.ServerID(server.ID)] = server
 }
 
 func (sl *ServerLookup) RemoveServer(server *ServerParts) {
diff --git a/dkron/server_lookup_test.go b/dkron/server_lookup_test.go
new file mode 100644
index 000000000..b60f112d9
--- /dev/null
+++ b/dkron/server_lookup_test.go
@@ -0,0 +1,73 @@
+// Copyright (c) HashiCorp, Inc.
+// SPDX-License-Identifier: BUSL-1.1
+
+package dkron
+
+import (
+	"testing"
+
+	"github.com/hashicorp/raft"
+	"github.com/stretchr/testify/require"
+)
+
+type testAddr struct {
+	addr string
+}
+
+func (ta *testAddr) Network() string {
+	return "tcp"
+}
+
+func (ta *testAddr) String() string {
+	return ta.addr
+}
+
+func TestAddServer(t *testing.T) {
+	// arrange
+	lookup := NewServerLookup()
+	id1, addr1 := "svr-1", "127.0.0.1:8300"
+	server1 := buildServerParts(id1, addr1)
+
+	// act
+	lookup.AddServer(server1)
+
+	// assert
+	servers := lookup.Servers()
+	require.EqualValuesf(t, 1, len(servers), "Expected 1 servers but got %v", len(servers))
+	require.Containsf(t, servers, server1, "Expected server %v to be in the list of servers", server1)
+
+	got, err := lookup.ServerAddr(raft.ServerID(id1))
+	require.NoErrorf(t, err, "Unexpected error: %v", err)
+	require.EqualValuesf(t, addr1, string(got), "Expected %v but got %v", addr1, got)
+
+	server := lookup.Server(raft.ServerAddress(addr1))
+	strAddr := server.RPCAddr.String()
+	require.EqualValuesf(t, addr1, strAddr, "Expected lookup to return address %v but got %v", addr1, strAddr)
+}
+
+func TestRemoveServer(t *testing.T) {
+	// arrange
+	lookup := NewServerLookup()
+	id1, addr1 := "svr-1", "127.0.0.1:8300"
+	server1 := buildServerParts(id1, addr1)
+	lookup.AddServer(server1)
+
+	// act
+	lookup.RemoveServer(server1)
+
+	// assert
+	servers := lookup.Servers()
+	require.EqualValuesf(t, 0, len(servers), "Expected 0 servers but got %v", len(servers))
+
+	require.Nilf(t, lookup.Server(raft.ServerAddress(addr1)), "Expected lookup to return nil")
+	addr, err := lookup.ServerAddr(raft.ServerID(id1))
+	require.Errorf(t, err, "Expected lookup to return error")
+	require.EqualValuesf(t, "", string(addr), "Expected empty address but got %v", addr)
+}
+
+func buildServerParts(id, addr string) *ServerParts {
+	return &ServerParts{
+		ID:      id,
+		RPCAddr: &testAddr{addr},
+	}
+}