From a771b82a7246ccab3d40dc0fe9d49ce683367ca1 Mon Sep 17 00:00:00 2001
From: Kenneth Jenkins <51246568+kenjenkins@users.noreply.github.com>
Date: Thu, 7 Dec 2023 12:21:10 -0800
Subject: [PATCH] storage/inmemory: fix Patch() error handling (#4838)

The Patch() method was intended to skip any records that do not
currently exist. However, currently inmemory.Backend.Patch() will return
ErrNotFound if the last record in the records slice is not found (it
will ignore any other previous records that are not found).

Update the error handling logic here to be consistent with the postgres
backend, and add a unit test to exercise this case.
---
 pkg/storage/inmemory/backend.go        |  4 ++--
 pkg/storage/storagetest/storagetest.go | 11 +++++++++++
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/pkg/storage/inmemory/backend.go b/pkg/storage/inmemory/backend.go
index 98357cda7..471235dad 100644
--- a/pkg/storage/inmemory/backend.go
+++ b/pkg/storage/inmemory/backend.go
@@ -269,12 +269,12 @@ func (backend *Backend) Patch(
 			// Skip any record that does not currently exist.
 			continue
 		} else if err != nil {
-			return
+			return serverVersion, patchedRecords, err
 		}
 		patchedRecords = append(patchedRecords, record)
 	}
 
-	return
+	return serverVersion, patchedRecords, nil
 }
 
 // patch updates the specified fields of an existing record, assuming the RWMutex is held.
diff --git a/pkg/storage/storagetest/storagetest.go b/pkg/storage/storagetest/storagetest.go
index a35d2b619..b6c1ee446 100644
--- a/pkg/storage/storagetest/storagetest.go
+++ b/pkg/storage/storagetest/storagetest.go
@@ -32,6 +32,17 @@ func TestBackendPatch(t *testing.T, ctx context.Context, backend storage.Backend
 		}
 	}
 
+	t.Run("not found", func(t *testing.T) {
+		mask, err := fieldmaskpb.New(&session.Session{}, "oauth_token")
+		require.NoError(t, err)
+
+		s := &session.Session{Id: "session-id-that-does-not-exist"}
+
+		_, updated, err := backend.Patch(ctx, []*databroker.Record{mkRecord(s)}, mask)
+		require.NoError(t, err)
+		assert.Empty(t, updated)
+	})
+
 	t.Run("basic", func(t *testing.T) {
 		// Populate an initial set of session records.
 		s1 := &session.Session{