Merge pull request #2960 from oasislabs/kostko/fix/rt-history-test

go/runtime/history: Fix flaky TestHistoryPrune test

go/runtime/history/history.go

@@ -140,6 +140,9 @@ func (h *runtimeHistory) LastConsensusHeight() (int64, error) {
 }
 
 func (h *runtimeHistory) GetBlock(ctx context.Context, round uint64) (*block.Block, error) {
+	if ctx.Err() != nil {
+		return nil, ctx.Err()
+	}
 	annBlk, err := h.db.getBlock(round)
 	if err != nil {
 		return nil, err
@@ -149,6 +152,9 @@ func (h *runtimeHistory) GetBlock(ctx context.Context, round uint64) (*block.Blo
 }
 
 func (h *runtimeHistory) GetLatestBlock(ctx context.Context) (*block.Block, error) {
+	if ctx.Err() != nil {
+		return nil, ctx.Err()
+	}
 	meta, err := h.db.metadata()
 	if err != nil {
 		return nil, err

go/runtime/history/history_test.go

@@ -185,6 +185,22 @@ func TestHistoryPrune(t *testing.T) {
 		t.Fatalf("failed to wait for prune to complete")
 	}
 
+	// Wait until the pruning transaction has been committed. This is needed because doneCh may be
+	// closed before the database transaction commits so the pruning may not yet be visible.
+	ctx, cancel := context.WithTimeout(context.Background(), recvTimeout)
+	defer cancel()
+	for {
+		_, err = history.GetBlock(ctx, 0)
+		if err == nil {
+			time.Sleep(10 * time.Millisecond)
+			continue
+		}
+
+		require.Error(err, "GetBlock should fail for pruned block 0")
+		require.Equal(roothash.ErrNotFound, err)
+		break
+	}
+
 	// Ensure we can only lookup the last 10 blocks.
 	for i := 0; i <= 50; i++ {
 		_, err = history.GetBlock(context.Background(), uint64(i))