storage: address TODO in TestPushTxnHeartbeatTimeout

Release note: None

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)

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 {