From 6115152f124fd14390864bb2b2b286e9e68b90bf Mon Sep 17 00:00:00 2001
From: Eric Schmidt <erschmid@google.com>
Date: Thu, 27 Oct 2022 12:21:50 -0700
Subject: [PATCH] chore(datastore): fix for Get & Multi tests (#6941)

* chore(datastore): fix for Get & Multi tests

* unit test fix
---
 datastore/datastore.go      |  6 ++-
 datastore/datastore_test.go | 92 ++++++++++++++++++-------------------
 2 files changed, 50 insertions(+), 48 deletions(-)

diff --git a/datastore/datastore.go b/datastore/datastore.go
index 2f5a1d609a58..e2afe17ee0ca 100644
--- a/datastore/datastore.go
+++ b/datastore/datastore.go
@@ -357,7 +357,8 @@ func (c *Client) Get(ctx context.Context, key *Key, dst interface{}) (err error)
 	if !c.readSettings.readTime.IsZero() {
 		opts = &pb.ReadOptions{
 			ConsistencyType: &pb.ReadOptions_ReadTime{
-				ReadTime: timestamppb.New(c.readSettings.readTime),
+				// Timestamp cannot be less than microseconds accuracy. See #6938
+				ReadTime: &timestamppb.Timestamp{Seconds: c.readSettings.readTime.Unix()},
 			},
 		}
 	}
@@ -389,7 +390,8 @@ func (c *Client) GetMulti(ctx context.Context, keys []*Key, dst interface{}) (er
 	if c.readSettings != nil && !c.readSettings.readTime.IsZero() {
 		opts = &pb.ReadOptions{
 			ConsistencyType: &pb.ReadOptions_ReadTime{
-				ReadTime: timestamppb.New(c.readSettings.readTime),
+				// Timestamp cannot be less than microseconds accuracy. See #6938
+				ReadTime: &timestamppb.Timestamp{Seconds: c.readSettings.readTime.Unix()},
 			},
 		}
 	}
diff --git a/datastore/datastore_test.go b/datastore/datastore_test.go
index 2671fa686bce..ba2ace4fdee9 100644
--- a/datastore/datastore_test.go
+++ b/datastore/datastore_test.go
@@ -17,7 +17,6 @@ package datastore
 import (
 	"context"
 	"errors"
-	"fmt"
 	"sort"
 	"strings"
 	"testing"
@@ -27,6 +26,7 @@ import (
 	"github.com/google/go-cmp/cmp"
 	pb "google.golang.org/genproto/googleapis/datastore/v1"
 	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/known/timestamppb"
 )
 
 func TestQueryConstruction(t *testing.T) {
@@ -306,29 +306,29 @@ func TestGetWithReadTime(t *testing.T) {
 			"A": {ValueType: &pb.Value_IntegerValue{IntegerValue: 1}},
 		},
 	}
-	fakeClient := &fakeDatastoreClient{
-		lookup: func(req *pb.LookupRequest) (*pb.LookupResponse, error) {
-			if !req.ReadOptions.GetReadTime().AsTime().Equal(tm) {
-				return nil, fmt.Errorf("read time mismatch: expected %v, got %v", tm,
-					req.ReadOptions.GetReadTime())
-			}
 
-			return &pb.LookupResponse{
-				Found: []*pb.EntityResult{
-					{
-						Entity:  e,
-						Version: 1,
-					},
-				},
-			}, nil
-		},
-	}
-
-	client := &Client{
-		client:       fakeClient,
-		readSettings: &readSettings{},
-	}
+	client, srv, cleanup := newMock(t)
+	defer cleanup()
 
+	srv.addRPC(&pb.LookupRequest{
+		ProjectId:  "projectID",
+		DatabaseId: "",
+		Keys: []*pb.Key{
+			keyToProto(k),
+		},
+		ReadOptions: &pb.ReadOptions{
+			ConsistencyType: &pb.ReadOptions_ReadTime{
+				ReadTime: &timestamppb.Timestamp{Seconds: tm.Unix()},
+			},
+		},
+	}, &pb.LookupResponse{
+		Found: []*pb.EntityResult{
+			{
+				Entity:  e,
+				Version: 1,
+			},
+		},
+	})
 	ctx := context.Background()
 	client.WithReadOptions(ReadTime(tm))
 	dst := &ent{}
@@ -362,32 +362,32 @@ func TestGetMultiWithReadTime(t *testing.T) {
 		},
 	}
 
-	fakeClient := &fakeDatastoreClient{
-		lookup: func(req *pb.LookupRequest) (*pb.LookupResponse, error) {
-
-			if !req.ReadOptions.GetReadTime().AsTime().Equal(tm) {
-				return nil, fmt.Errorf("read time mismatch: expected %v, got %v", tm,
-					req.ReadOptions.GetReadTime())
-			}
+	client, srv, cleanup := newMock(t)
+	defer cleanup()
 
-			return &pb.LookupResponse{
-				Found: []*pb.EntityResult{
-					{
-						Entity:  e,
-						Version: 1,
-					}, {
-						Entity:  e2,
-						Version: 1,
-					},
-				},
-			}, nil
+	srv.addRPC(&pb.LookupRequest{
+		ProjectId:  "projectID",
+		DatabaseId: "",
+		Keys: []*pb.Key{
+			keyToProto(k[0]),
+			keyToProto(k[1]),
 		},
-	}
-
-	client := &Client{
-		client:       fakeClient,
-		readSettings: &readSettings{},
-	}
+		ReadOptions: &pb.ReadOptions{
+			ConsistencyType: &pb.ReadOptions_ReadTime{
+				ReadTime: &timestamppb.Timestamp{Seconds: tm.Unix()},
+			},
+		},
+	}, &pb.LookupResponse{
+		Found: []*pb.EntityResult{
+			{
+				Entity:  e,
+				Version: 1,
+			}, {
+				Entity:  e2,
+				Version: 1,
+			},
+		},
+	})
 
 	ctx := context.Background()
 	client.WithReadOptions(ReadTime(tm))