Rename mock router to memory router and add tests

pkg/registry/registry_test.go

@@ -48,7 +48,7 @@ func TestMirrorHandler(t *testing.T) {
 		"first-peer-error":  {unreachableAddrPort, goodAddrPort},
 		"last-peer-working": {badAddrPort, badAddrPort, goodAddrPort},
 	}
-	router := routing.NewMockRouter(resolver, netip.AddrPort{})
+	router := routing.NewMemoryRouter(resolver, netip.AddrPort{})
 	reg := NewRegistry(nil, router)
 
 	tests := []struct {

pkg/routing/mock.go → pkg/routing/memory.go

@@ -6,26 +6,26 @@ import (
 	"sync"
 )
 
-type MockRouter struct {
+type MemoryRouter struct {
 	resolver map[string][]netip.AddrPort
 	self     netip.AddrPort
 	mx       sync.RWMutex
 }
 
-func NewMockRouter(resolver map[string][]netip.AddrPort, self netip.AddrPort) *MockRouter {
-	return &MockRouter{
+func NewMemoryRouter(resolver map[string][]netip.AddrPort, self netip.AddrPort) *MemoryRouter {
+	return &MemoryRouter{
 		resolver: resolver,
 		self:     self,
 	}
 }
 
-func (m *MockRouter) Ready(ctx context.Context) (bool, error) {
+func (m *MemoryRouter) Ready(ctx context.Context) (bool, error) {
 	m.mx.RLock()
 	defer m.mx.RUnlock()
 	return len(m.resolver) > 0, nil
 }
 
-func (m *MockRouter) Resolve(ctx context.Context, key string, allowSelf bool, count int) (<-chan netip.AddrPort, error) {
+func (m *MemoryRouter) Resolve(ctx context.Context, key string, allowSelf bool, count int) (<-chan netip.AddrPort, error) {
 	peerCh := make(chan netip.AddrPort, count)
 	peers, ok := m.resolver[key]
 	// If not peers exist close the channel to stop any consumer.
@@ -34,8 +34,6 @@ func (m *MockRouter) Resolve(ctx context.Context, key string, allowSelf bool, co
 		return peerCh, nil
 	}
 	go func() {
-		m.mx.RLock()
-		defer m.mx.RUnlock()
 		for _, peer := range peers {
 			peerCh <- peer
 		}
@@ -44,16 +42,24 @@ func (m *MockRouter) Resolve(ctx context.Context, key string, allowSelf bool, co
 	return peerCh, nil
 }
 
-func (m *MockRouter) Advertise(ctx context.Context, keys []string) error {
-	m.mx.Lock()
-	defer m.mx.Unlock()
+func (m *MemoryRouter) Advertise(ctx context.Context, keys []string) error {
 	for _, key := range keys {
-		m.resolver[key] = []netip.AddrPort{m.self}
+		m.Add(key, m.self)
 	}
 	return nil
 }
 
-func (m *MockRouter) LookupKey(key string) ([]netip.AddrPort, bool) {
+func (m *MemoryRouter) Add(key string, ap netip.AddrPort) {
+	m.mx.Lock()
+	defer m.mx.Unlock()
+	if v, ok := m.resolver[key]; ok {
+		m.resolver[key] = append(v, ap)
+		return
+	}
+	m.resolver[key] = []netip.AddrPort{ap}
+}
+
+func (m *MemoryRouter) Lookup(key string) ([]netip.AddrPort, bool) {
 	m.mx.RLock()
 	defer m.mx.RUnlock()
 	v, ok := m.resolver[key]

pkg/routing/memory_test.go

@@ -0,0 +1,33 @@
+package routing
+
+import (
+	"context"
+	"net/netip"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestMemoryRouter(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	r := NewMemoryRouter(map[string][]netip.AddrPort{}, netip.AddrPort{})
+
+	isReady, err := r.Ready(ctx)
+	require.NoError(t, err)
+	require.False(t, isReady)
+	r.Advertise(ctx, []string{"foo"})
+	isReady, err = r.Ready(ctx)
+	require.NoError(t, err)
+	require.True(t, isReady)
+
+	r.Add("foo", netip.MustParseAddrPort("127.0.0.1:9090"))
+	peerCh, err := r.Resolve(ctx, "foo", true, 2)
+	require.NoError(t, err)
+	peers := []netip.AddrPort{}
+	for peer := range peerCh {
+		peers = append(peers, peer)
+	}
+	require.Len(t, peers, 2)
+}

pkg/state/state_test.go

@@ -46,7 +46,7 @@ func TestBasic(t *testing.T) {
 				imgs = append(imgs, img)
 			}
 			ociClient := oci.NewMockClient(imgs)
-			router := routing.NewMockRouter(map[string][]netip.AddrPort{}, netip.MustParseAddrPort("127.0.0.1:5000"))
+			router := routing.NewMemoryRouter(map[string][]netip.AddrPort{}, netip.MustParseAddrPort("127.0.0.1:5000"))
 
 			ctx, cancel := context.WithCancel(context.TODO())
 			go func() {
@@ -57,14 +57,14 @@ func TestBasic(t *testing.T) {
 			require.NoError(t, err)
 
 			for _, img := range imgs {
-				peers, ok := router.LookupKey(img.Digest.String())
+				peers, ok := router.Lookup(img.Digest.String())
 				require.True(t, ok)
 				require.Len(t, peers, 1)
 				tagName, ok := img.TagName()
 				if !ok {
 					continue
 				}
-				peers, ok = router.LookupKey(tagName)
+				peers, ok = router.Lookup(tagName)
 				if img.IsLatestTag() && !tt.resolveLatestTag {
 					require.False(t, ok)
 					continue