github.com/filecoin-project/lassie@v0.23.0/pkg/internal/testutil/verifier.go (about)

     1  package testutil
     2  
     3  import (
     4  	"context"
     5  	"sort"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/benbjohnson/clock"
    10  	"github.com/filecoin-project/lassie/pkg/events"
    11  	"github.com/filecoin-project/lassie/pkg/types"
    12  	blocks "github.com/ipfs/go-block-format"
    13  	"github.com/ipfs/go-cid"
    14  	"github.com/libp2p/go-libp2p/core/peer"
    15  	"github.com/multiformats/go-multicodec"
    16  	"github.com/stretchr/testify/require"
    17  )
    18  
    19  type ExpectedActionsAtTime struct {
    20  	AfterStart           time.Duration
    21  	ReceivedConnections  []peer.ID
    22  	ReceivedRetrievals   []peer.ID
    23  	ServedRetrievals     []RemoteStats
    24  	CompletedRetrievals  []peer.ID
    25  	CandidatesDiscovered []DiscoveredCandidate
    26  	ExpectedEvents       []types.RetrievalEvent
    27  	ExpectedMetrics      []SessionMetric
    28  }
    29  
    30  type RemoteStats struct {
    31  	Peer      peer.ID
    32  	Root      cid.Cid
    33  	ByteCount uint64
    34  	Blocks    []cid.Cid
    35  	Err       struct{}
    36  }
    37  
    38  type RetrievalVerifier struct {
    39  	ExpectedSequence []ExpectedActionsAtTime
    40  }
    41  
    42  type RunRetrieval func(cb func(types.RetrievalEvent)) (*types.RetrievalStats, error)
    43  
    44  type VerifierClient interface {
    45  	VerifyConnectionsReceived(ctx context.Context, t *testing.T, afterStart time.Duration, expectedConnections []peer.ID)
    46  	VerifyRetrievalsReceived(ctx context.Context, t *testing.T, afterStart time.Duration, expectedRetrievals []peer.ID)
    47  	VerifyRetrievalsServed(ctx context.Context, t *testing.T, afterStart time.Duration, expectedServed []RemoteStats)
    48  	VerifyRetrievalsCompleted(ctx context.Context, t *testing.T, afterStart time.Duration, expectedRetrievals []peer.ID)
    49  }
    50  
    51  func (rv RetrievalVerifier) RunWithVerification(
    52  	ctx context.Context,
    53  	t *testing.T,
    54  	clock *clock.Mock,
    55  	client VerifierClient,
    56  	mockCandidateSource *MockCandidateSource,
    57  	mockSession *MockSession,
    58  	cancelFunc context.CancelFunc,
    59  	cancelAfter time.Duration,
    60  	runRetrievals []RunRetrieval,
    61  ) []types.RetrievalResult {
    62  
    63  	resultChan := make(chan types.RetrievalResult, len(runRetrievals))
    64  	asyncCollectingEventsListener := NewAsyncCollectingEventsListener(ctx)
    65  	for _, runRetrieval := range runRetrievals {
    66  		runRetrieval := runRetrieval
    67  		go func() {
    68  			result, err := runRetrieval(asyncCollectingEventsListener.Collect)
    69  			resultChan <- types.RetrievalResult{Stats: result, Err: err}
    70  		}()
    71  	}
    72  	currentTime := time.Duration(0)
    73  	for _, expectedActionsAtTime := range rv.ExpectedSequence {
    74  		if cancelAfter != 0 && cancelFunc != nil && expectedActionsAtTime.AfterStart >= cancelAfter {
    75  			t.Logf("Cancelling context @ %s", cancelAfter)
    76  			cancelFunc()
    77  			cancelFunc = nil
    78  		}
    79  		clock.Add(expectedActionsAtTime.AfterStart - currentTime)
    80  		currentTime = expectedActionsAtTime.AfterStart
    81  		t.Logf("ExpectedAction.AfterStart=%s", expectedActionsAtTime.AfterStart)
    82  		asyncCollectingEventsListener.VerifyNextEvents(t, expectedActionsAtTime.AfterStart, expectedActionsAtTime.ExpectedEvents)
    83  		if mockSession != nil {
    84  			mockSession.VerifyMetricsAt(ctx, t, expectedActionsAtTime.AfterStart, expectedActionsAtTime.ExpectedMetrics)
    85  		}
    86  		if mockCandidateSource != nil {
    87  			mockCandidateSource.VerifyCandidatesDiscovered(ctx, t, expectedActionsAtTime.AfterStart, expectedActionsAtTime.CandidatesDiscovered)
    88  		}
    89  		if client != nil {
    90  			client.VerifyConnectionsReceived(ctx, t, expectedActionsAtTime.AfterStart, expectedActionsAtTime.ReceivedConnections)
    91  			client.VerifyRetrievalsReceived(ctx, t, expectedActionsAtTime.AfterStart, expectedActionsAtTime.ReceivedRetrievals)
    92  			client.VerifyRetrievalsServed(ctx, t, expectedActionsAtTime.AfterStart, expectedActionsAtTime.ServedRetrievals)
    93  			client.VerifyRetrievalsCompleted(ctx, t, expectedActionsAtTime.AfterStart, expectedActionsAtTime.CompletedRetrievals)
    94  		}
    95  	}
    96  	results := make([]types.RetrievalResult, 0, len(runRetrievals))
    97  	for i := 0; i < len(runRetrievals); i++ {
    98  		select {
    99  		case result := <-resultChan:
   100  			results = append(results, result)
   101  		case <-ctx.Done():
   102  			require.FailNowf(t, "did not complete retrieval", "got %d of %d", i, len(runRetrievals))
   103  		}
   104  	}
   105  	return results
   106  }
   107  
   108  func BlockReceivedActions(baseTime time.Time, baseAfterStart time.Duration, rid types.RetrievalID, candidate types.RetrievalCandidate, protocol multicodec.Code, blockTime time.Duration, blks []blocks.Block) []ExpectedActionsAtTime {
   109  	var actions []ExpectedActionsAtTime
   110  	for i, blk := range blks {
   111  		actions = append(actions, ExpectedActionsAtTime{
   112  			AfterStart: baseAfterStart + time.Duration(i)*blockTime,
   113  			ExpectedEvents: []types.RetrievalEvent{
   114  				events.BlockReceived(baseTime.Add(baseAfterStart+time.Duration(i)*blockTime), rid, candidate, protocol, uint64(len(blk.RawData()))),
   115  			},
   116  		})
   117  	}
   118  	return actions
   119  }
   120  
   121  func SortActions(actions []ExpectedActionsAtTime) []ExpectedActionsAtTime {
   122  	sort.Slice(actions, func(i, j int) bool {
   123  		return actions[i].AfterStart < actions[j].AfterStart
   124  	})
   125  	return actions
   126  }