From aa97806eede50a55959358b80f5dbadfeb90e3c1 Mon Sep 17 00:00:00 2001
From: Nathan VanBenschoten <nvanbenschoten@gmail.com>
Date: Tue, 16 Jul 2019 19:41:20 -0400
Subject: [PATCH] storage: address TODO in TestPushTxnHeartbeatTimeout

Release note: None
---
 pkg/storage/replica_test.go                  | 15 ++++++---------
 pkg/storage/txn_recovery_integration_test.go |  4 ++--
 2 files changed, 8 insertions(+), 11 deletions(-)

diff --git a/pkg/storage/replica_test.go b/pkg/storage/replica_test.go
index a05d8b8593cf..d534b852e75c 100644
--- a/pkg/storage/replica_test.go
+++ b/pkg/storage/replica_test.go
@@ -5159,15 +5159,12 @@ func TestPushTxnHeartbeatTimeout(t *testing.T) {
 				t.Fatalf("%d: %s", i, pErr)
 			}
 		case roachpb.STAGING:
-			// TODO(nvanbenschoten): Avoid writing directly to the engine once
-			// there's a way to create a STAGING transaction record.
-			txnKey := keys.TransactionKey(pushee.Key, pushee.ID)
-			txnRecord := pushee.AsRecord()
-			txnRecord.Status = roachpb.STAGING
-			if err := engine.MVCCPutProto(
-				context.Background(), tc.repl.store.Engine(), nil, txnKey, hlc.Timestamp{}, nil, &txnRecord,
-			); err != nil {
-				t.Fatal(err)
+			et, etH := endTxnArgs(pushee, true)
+			et.InFlightWrites = []roachpb.SequencedWrite{
+				{Key: key, Sequence: 1},
+			}
+			if _, pErr := client.SendWrappedWith(context.Background(), tc.Sender(), etH, &et); pErr != nil {
+				t.Fatalf("%d: %s", i, pErr)
 			}
 		default:
 			t.Fatalf("unexpected status: %v", test.status)
diff --git a/pkg/storage/txn_recovery_integration_test.go b/pkg/storage/txn_recovery_integration_test.go
index a064e0cefc13..41256189087a 100644
--- a/pkg/storage/txn_recovery_integration_test.go
+++ b/pkg/storage/txn_recovery_integration_test.go
@@ -71,8 +71,8 @@ func TestTxnRecoveryFromStaging(t *testing.T) {
 		// state. Include both writes as the EndTransaction's in-flight writes.
 		et, etH := endTxnArgs(txn, true)
 		et.InFlightWrites = []roachpb.SequencedWrite{
-			{Key: keyA, Sequence: 0},
-			{Key: keyB, Sequence: 1},
+			{Key: keyA, Sequence: 1},
+			{Key: keyB, Sequence: 2},
 		}
 		etReply, pErr := client.SendWrappedWith(ctx, store.TestSender(), etH, &et)
 		if pErr != nil {