From 5e8179c6fb8663bc51557d453fa1b7b82d8f0940 Mon Sep 17 00:00:00 2001
From: Michael Gaffney <mgaffney@users.noreply.github.com>
Date: Thu, 30 Jul 2020 13:44:14 -0400
Subject: [PATCH] Allow lookupAfterWrite for resources with Private Ids (#229)

---
 internal/db/read_writer.go      |  9 ++--
 internal/db/read_writer_test.go | 89 +++++++++++++++++++++++++++++++++
 2 files changed, 92 insertions(+), 6 deletions(-)

diff --git a/internal/db/read_writer.go b/internal/db/read_writer.go
index 6680f98eb..53e08abf0 100644
--- a/internal/db/read_writer.go
+++ b/internal/db/read_writer.go
@@ -205,13 +205,10 @@ func (rw *Db) lookupAfterWrite(ctx context.Context, i interface{}, opt ...Option
 	if !withLookup {
 		return nil
 	}
-	if _, ok := i.(ResourcePublicIder); ok {
-		if err := rw.LookupByPublicId(ctx, i.(ResourcePublicIder), opt...); err != nil {
-			return fmt.Errorf("lookup after write: %w", err)
-		}
-		return nil
+	if err := rw.LookupById(ctx, i, opt...); err != nil {
+		return fmt.Errorf("lookup after write: %w", err)
 	}
-	return errors.New("lookup after write: not a resource with an id")
+	return nil
 }
 
 // Create an object in the db with options: WithOplog, NewOplogMsg and
diff --git a/internal/db/read_writer_test.go b/internal/db/read_writer_test.go
index 289ac0891..db07b96f9 100644
--- a/internal/db/read_writer_test.go
+++ b/internal/db/read_writer_test.go
@@ -2788,3 +2788,92 @@ func TestDb_oplogMsgsForItems(t *testing.T) {
 		})
 	}
 }
+
+func TestDb_lookupAfterWrite(t *testing.T) {
+	t.Parallel()
+	db, _ := TestSetup(t, "postgres")
+	scooter := testScooter(t, db, "", 0)
+	user := testUser(t, db, "", "", "")
+	type args struct {
+		resourceWithIder interface{}
+		opt              []Option
+	}
+	tests := []struct {
+		name       string
+		underlying *gorm.DB
+		args       args
+		wantErr    bool
+		want       proto.Message
+		wantIsErr  error
+	}{
+		{
+			name:       "simple-private-id",
+			underlying: db,
+			args: args{
+				resourceWithIder: scooter,
+				opt:              []Option{WithLookup(true)},
+			},
+			wantErr: false,
+			want:    scooter,
+		},
+		{
+			name:       "simple-public-id",
+			underlying: db,
+			args: args{
+				resourceWithIder: user,
+				opt:              []Option{WithLookup(true)},
+			},
+			wantErr: false,
+			want:    user,
+		},
+		{
+			name:       "no-lookup-private-id",
+			underlying: db,
+			args: args{
+				resourceWithIder: scooter,
+				opt:              []Option{WithLookup(false)},
+			},
+			wantErr: false,
+			want:    nil,
+		},
+		{
+			name:       "no-lookup-public-id",
+			underlying: db,
+			args: args{
+				resourceWithIder: user,
+				opt:              []Option{WithLookup(false)},
+			},
+			wantErr: false,
+			want:    nil,
+		},
+		{
+			name:       "not-an-ider",
+			underlying: db,
+			args: args{
+				resourceWithIder: &db_test.NotIder{},
+				opt:              []Option{WithLookup(true)},
+			},
+			wantErr: true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert, require := assert.New(t), require.New(t)
+			rw := &Db{
+				underlying: tt.underlying,
+			}
+			cloner, ok := tt.args.resourceWithIder.(db_test.Cloner)
+			require.True(ok)
+			cp := cloner.Clone()
+			err := rw.lookupAfterWrite(context.Background(), cp, tt.args.opt...)
+			if tt.wantErr {
+				require.Error(err)
+				return
+			}
+			require.NoError(err)
+			if tt.want != nil {
+				assert.True(proto.Equal(tt.want, cp.(proto.Message)))
+			}
+		})
+	}
+}