From 8d592f4adc4558b0761be5deaed15c17cd43ca4c Mon Sep 17 00:00:00 2001
From: Terence Tsao <terence@prysmaticlabs.com>
Date: Mon, 4 Nov 2019 11:25:55 -0800
Subject: [PATCH] Unskip minimal spec test for finalization

---
 ...ification_and_finalization_minimal_test.go |  1 -
 .../justification_and_finalization_test.go    | 43 +++++++++++++++++--
 2 files changed, 39 insertions(+), 5 deletions(-)

diff --git a/beacon-chain/core/epoch/spectest/justification_and_finalization_minimal_test.go b/beacon-chain/core/epoch/spectest/justification_and_finalization_minimal_test.go
index 992821791f6f..57775ffc130a 100644
--- a/beacon-chain/core/epoch/spectest/justification_and_finalization_minimal_test.go
+++ b/beacon-chain/core/epoch/spectest/justification_and_finalization_minimal_test.go
@@ -5,6 +5,5 @@ import (
 )
 
 func TestJustificationAndFinalizationMinimal(t *testing.T) {
-	t.Skip("Fails for could not get target atts current epoch")
 	runJustificationAndFinalizationTests(t, "minimal")
 }
diff --git a/beacon-chain/core/epoch/spectest/justification_and_finalization_test.go b/beacon-chain/core/epoch/spectest/justification_and_finalization_test.go
index 2df30f929b50..d09f080ae5c3 100644
--- a/beacon-chain/core/epoch/spectest/justification_and_finalization_test.go
+++ b/beacon-chain/core/epoch/spectest/justification_and_finalization_test.go
@@ -1,7 +1,9 @@
 package spectest
 
 import (
+	"bytes"
 	"context"
+	"fmt"
 	"path"
 	"testing"
 
@@ -32,19 +34,19 @@ func runJustificationAndFinalizationTests(t *testing.T, config string) {
 // This is a subset of state.ProcessEpoch. The spec test defines input data for
 // `justification_and_finalization` only.
 func processJustificationAndFinalizationWrapper(t *testing.T, state *pb.BeaconState) (*pb.BeaconState, error) {
-	prevEpochAtts, err := epoch.MatchAttestations(state, helpers.PrevEpoch(state))
+	prevEpochAtts, err := targetAtts(state, helpers.PrevEpoch(state))
 	if err != nil {
 		t.Fatalf("could not get target atts prev epoch %d: %v", helpers.PrevEpoch(state), err)
 	}
-	currentEpochAtts, err := epoch.MatchAttestations(state, helpers.CurrentEpoch(state))
+	currentEpochAtts, err := targetAtts(state, helpers.CurrentEpoch(state))
 	if err != nil {
 		t.Fatalf("could not get target atts current epoch %d: %v", helpers.CurrentEpoch(state), err)
 	}
-	prevEpochAttestedBalance, err := epoch.AttestingBalance(state, prevEpochAtts.Target)
+	prevEpochAttestedBalance, err := epoch.AttestingBalance(state, prevEpochAtts)
 	if err != nil {
 		t.Fatalf("could not get attesting balance prev epoch: %v", err)
 	}
-	currentEpochAttestedBalance, err := epoch.AttestingBalance(state, currentEpochAtts.Target)
+	currentEpochAttestedBalance, err := epoch.AttestingBalance(state, currentEpochAtts)
 	if err != nil {
 		t.Fatalf("could not get attesting balance current epoch: %v", err)
 	}
@@ -72,3 +74,36 @@ func processJustificationAndFinalizationPrecomputeWrapper(t *testing.T, state *p
 
 	return state, nil
 }
+
+func targetAtts(state *pb.BeaconState, epoch uint64) ([]*pb.PendingAttestation, error) {
+	currentEpoch := helpers.CurrentEpoch(state)
+	previousEpoch := helpers.PrevEpoch(state)
+
+	// Input epoch for matching the source attestations has to be within range
+	// of current epoch & previous epoch.
+	if epoch != currentEpoch && epoch != previousEpoch {
+		return nil, fmt.Errorf("input epoch: %d != current epoch: %d or previous epoch: %d",
+			epoch, currentEpoch, previousEpoch)
+	}
+
+	// Decide if the source attestations are coming from current or previous epoch.
+	var srcAtts []*pb.PendingAttestation
+	if epoch == currentEpoch {
+		srcAtts = state.CurrentEpochAttestations
+	} else {
+		srcAtts = state.PreviousEpochAttestations
+	}
+	targetRoot, err := helpers.BlockRoot(state, epoch)
+	if err != nil {
+		return nil, err
+	}
+
+	tgtAtts := make([]*pb.PendingAttestation, 0, len(srcAtts))
+	for _, srcAtt := range srcAtts {
+		if bytes.Equal(srcAtt.Data.Target.Root, targetRoot) {
+			tgtAtts = append(tgtAtts, srcAtt)
+		}
+	}
+
+	return tgtAtts, nil
+}