pass time.Time channel to PopWithTimer, move test

dot/state/transaction.go

@@ -50,8 +50,8 @@ func (s *TransactionState) Pop() *transaction.ValidTransaction {
 
 // PopWithTimer returns the next valid transaction from the queue.
 // When the timer expires, it returns `nil`.
-func (s *TransactionState) PopWithTimer(timer *time.Timer) (transaction *transaction.ValidTransaction) {
-	return s.queue.PopWithTimer(timer)
+func (s *TransactionState) PopWithTimer(timerCh <-chan time.Time) (transaction *transaction.ValidTransaction) {
+	return s.queue.PopWithTimer(timerCh)
 }
 
 // Peek returns the head of the queue without removing it

lib/babe/build.go

@@ -179,7 +179,7 @@ func (b *BlockBuilder) buildBlockExtrinsics(slot Slot, rt runtime.Instance) []*t
 	slotTimer := time.NewTimer(timeout)
 
 	for {
-		txn := b.transactionState.PopWithTimer(slotTimer)
+		txn := b.transactionState.PopWithTimer(slotTimer.C)
 		slotTimerExpired := txn == nil
 		if slotTimerExpired {
 			break

lib/babe/state.go

@@ -56,7 +56,7 @@ type TransactionState interface {
 	Push(vt *transaction.ValidTransaction) (common.Hash, error)
 	Pop() *transaction.ValidTransaction
 	Peek() *transaction.ValidTransaction
-	PopWithTimer(timer *time.Timer) (tx *transaction.ValidTransaction)
+	PopWithTimer(timerCh <-chan time.Time) (tx *transaction.ValidTransaction)
 }
 
 // EpochState is the interface for epoch methods

lib/transaction/priority_queue.go

@@ -143,7 +143,7 @@ func (spq *PriorityQueue) Push(txn *ValidTransaction) (common.Hash, error) {
 
 // PopWithTimer returns the next valid transaction from the queue.
 // When the timer expires, it returns `nil`.
-func (spq *PriorityQueue) PopWithTimer(timer *time.Timer) (transaction *ValidTransaction) {
+func (spq *PriorityQueue) PopWithTimer(timerCh <-chan time.Time) (transaction *ValidTransaction) {
 	transaction = spq.Pop()
 	if transaction != nil {
 		return transaction
@@ -156,7 +156,7 @@ func (spq *PriorityQueue) PopWithTimer(timer *time.Timer) (transaction *ValidTra
 
 		for {
 			select {
-			case <-timer.C:
+			case <-timerCh:
 				transactionChannel <- nil
 				return
 			case <-pollTicker.C:

lib/transaction/priority_queue_integration_test.go

@@ -0,0 +1,57 @@
+// Copyright 2021 ChainSafe Systems (ON)
+// SPDX-License-Identifier: LGPL-3.0-only
+
+//go:build integration
+
+package transaction
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_PopWithTimer(t *testing.T) {
+	pq := NewPriorityQueue()
+	slotTimer := time.NewTimer(time.Second)
+
+	tests := []*ValidTransaction{
+		{
+			Extrinsic: []byte("a"),
+			Validity:  &Validity{Priority: 1},
+		},
+		{
+			Extrinsic: []byte("b"),
+			Validity:  &Validity{Priority: 4},
+		},
+		{
+			Extrinsic: []byte("c"),
+			Validity:  &Validity{Priority: 2},
+		},
+		{
+			Extrinsic: []byte("d"),
+			Validity:  &Validity{Priority: 17},
+		},
+		{
+			Extrinsic: []byte("e"),
+			Validity:  &Validity{Priority: 2},
+		},
+	}
+
+	expected := []int{3, 1, 2, 4, 0}
+
+	for _, test := range tests {
+		pq.Push(test)
+	}
+
+	counter := 0
+	for {
+		txn := pq.PopWithTimer(slotTimer.C)
+		if txn == nil {
+			break
+		}
+		assert.Equal(t, tests[expected[counter]], txn)
+		counter++
+	}
+}

lib/transaction/priority_queue_test.go

@@ -273,50 +273,6 @@ func TestRemoveExtrinsic(t *testing.T) {
 	}
 }
 
-func Test_PopWithTimer(t *testing.T) {
-	pq := NewPriorityQueue()
-	slotTimer := time.NewTimer(time.Second)
-
-	tests := []*ValidTransaction{
-		{
-			Extrinsic: []byte("a"),
-			Validity:  &Validity{Priority: 1},
-		},
-		{
-			Extrinsic: []byte("b"),
-			Validity:  &Validity{Priority: 4},
-		},
-		{
-			Extrinsic: []byte("c"),
-			Validity:  &Validity{Priority: 2},
-		},
-		{
-			Extrinsic: []byte("d"),
-			Validity:  &Validity{Priority: 17},
-		},
-		{
-			Extrinsic: []byte("e"),
-			Validity:  &Validity{Priority: 2},
-		},
-	}
-
-	expected := []int{3, 1, 2, 4, 0}
-
-	for _, test := range tests {
-		pq.Push(test)
-	}
-
-	counter := 0
-	for {
-		txn := pq.PopWithTimer(slotTimer)
-		if txn == nil {
-			break
-		}
-		assert.Equal(t, tests[expected[counter]], txn)
-		counter++
-	}
-}
-
 func Test_PriorityQueue_PopWithTimer(t *testing.T) {
 	t.Parallel()
 
@@ -382,7 +338,7 @@ func Test_PriorityQueue_PopWithTimer(t *testing.T) {
 				close(modifyDone)
 			}
 
-			transaction := queue.PopWithTimer(testCase.timer)
+			transaction := queue.PopWithTimer(testCase.timer.C)
 			<-modifyDone
 			testCase.timer.Stop()
 			assert.Equal(t, testCase.transaction, transaction)