handle nil in mocks

Layr-Labs · Dec 11, 2024 · 7edad7a · 7edad7a
1 parent 844817e
commit 7edad7a
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/core/mock/writer.go b/core/mock/writer.go
@@ -265,12 +265,18 @@ func (t *MockWriter) GetNumRelays(ctx context.Context) (uint32, error) {
 
 func (t *MockWriter) GetRelayURL(ctx context.Context, key uint32) (string, error) {
 	args := t.Called()
+	if args.Get(0) == nil {
+		return "", args.Error(1)
+	}
 	result := args.Get(0)
 	return result.(string), args.Error(1)
 }
 
 func (t *MockWriter) GetRelayURLs(ctx context.Context) (map[uint32]string, error) {
 	args := t.Called()
+	if args.Get(0) == nil {
+		return nil, args.Error(1)
+	}
 	result := args.Get(0)
 	if result == nil {
 		return nil, args.Error(1)