1 files changed, 42 insertions, 4 deletions

diff --git a/pkg/p9/messages_test.go b/pkg/p9/messages_test.go
index 10a0587cf..513b30e8b 100644
--- a/pkg/p9/messages_test.go
+++ b/pkg/p9/messages_test.go
@@ -399,8 +399,9 @@ func TestEncodeDecode(t *testing.T) {
 }
 
 func TestMessageStrings(t *testing.T) {
-	for typ, fn := range messageRegistry {
-		if fn != nil {
+	for typ := range msgRegistry.factories {
+		entry := &msgRegistry.factories[typ]
+		if entry.create != nil {
 			name := fmt.Sprintf("%+v", typ)
 			t.Run(name, func(t *testing.T) {
 				defer func() { // Ensure no panic.
@@ -408,7 +409,7 @@ func TestMessageStrings(t *testing.T) {
 						t.Errorf("printing %s failed: %v", name, r)
 					}
 				}()
-				m := fn()
+				m := entry.create()
 				_ = fmt.Sprintf("%v", m)
 				err := ErrInvalidMsgType{MsgType(typ)}
 				_ = err.Error()
@@ -426,5 +427,42 @@ func TestRegisterDuplicate(t *testing.T) {
 	}()
 
 	// Register a duplicate.
-	register(MsgRlerror, func() message { return &Rlerror{} })
+	msgRegistry.register(MsgRlerror, func() message { return &Rlerror{} })
+}
+
+func TestMsgCache(t *testing.T) {
+	// Cache starts empty.
+	if got, want := len(msgRegistry.factories[MsgRlerror].cache), 0; got != want {
+		t.Errorf("Wrong cache size, got: %d, want: %d", got, want)
+	}
+
+	// Message can be created with an empty cache.
+	msg, err := msgRegistry.get(0, MsgRlerror)
+	if err != nil {
+		t.Errorf("msgRegistry.get(): %v", err)
+	}
+	if got, want := len(msgRegistry.factories[MsgRlerror].cache), 0; got != want {
+		t.Errorf("Wrong cache size, got: %d, want: %d", got, want)
+	}
+
+	// Check that message is added to the cache when returned.
+	msgRegistry.put(msg)
+	if got, want := len(msgRegistry.factories[MsgRlerror].cache), 1; got != want {
+		t.Errorf("Wrong cache size, got: %d, want: %d", got, want)
+	}
+
+	// Check that returned message is reused.
+	if got, err := msgRegistry.get(0, MsgRlerror); err != nil {
+		t.Errorf("msgRegistry.get(): %v", err)
+	} else if msg != got {
+		t.Errorf("Message not reused, got: %d, want: %d", got, msg)
+	}
+
+	// Check that cache doesn't grow beyond max size.
+	for i := 0; i < maxCacheSize+1; i++ {
+		msgRegistry.put(&Rlerror{})
+	}
+	if got, want := len(msgRegistry.factories[MsgRlerror].cache), maxCacheSize; got != want {
+		t.Errorf("Wrong cache size, got: %d, want: %d", got, want)
+	}
 }