From ea69386519519b3c3b309467b96d2e6dd9bd2e76 Mon Sep 17 00:00:00 2001
From: Jakub Sztandera <kubuxu@protocol.ai>
Date: Fri, 23 Oct 2020 20:46:06 +0200
Subject: [PATCH] Fix flaky TestTimedBSSimple

Signed-off-by: Jakub Sztandera <kubuxu@protocol.ai>
---
 lib/timedbs/timedbs.go      | 14 ++++++++++----
 lib/timedbs/timedbs_test.go | 23 +++++++++++++++--------
 2 files changed, 25 insertions(+), 12 deletions(-)

diff --git a/lib/timedbs/timedbs.go b/lib/timedbs/timedbs.go
index bb03a59e9..c5c1a8fe0 100644
--- a/lib/timedbs/timedbs.go
+++ b/lib/timedbs/timedbs.go
@@ -8,6 +8,7 @@ import (
 
 	blocks "github.com/ipfs/go-block-format"
 	"github.com/ipfs/go-cid"
+	"github.com/raulk/clock"
 	"go.uber.org/multierr"
 
 	"github.com/filecoin-project/lotus/build"
@@ -24,9 +25,10 @@ import (
 type TimedCacheBS struct {
 	mu               sync.RWMutex
 	active, inactive blockstore.MemStore
-
-	interval time.Duration
-	closeCh  chan struct{}
+	clock            clock.Clock
+	interval         time.Duration
+	closeCh          chan struct{}
+	doneRotatingCh   chan struct{}
 }
 
 func NewTimedCacheBS(cacheTime time.Duration) *TimedCacheBS {
@@ -34,6 +36,7 @@ func NewTimedCacheBS(cacheTime time.Duration) *TimedCacheBS {
 		active:   blockstore.NewTemporary(),
 		inactive: blockstore.NewTemporary(),
 		interval: cacheTime,
+		clock:    build.Clock,
 	}
 }
 
@@ -45,12 +48,15 @@ func (t *TimedCacheBS) Start(ctx context.Context) error {
 	}
 	t.closeCh = make(chan struct{})
 	go func() {
-		ticker := build.Clock.Ticker(t.interval)
+		ticker := t.clock.Ticker(t.interval)
 		defer ticker.Stop()
 		for {
 			select {
 			case <-ticker.C:
 				t.rotate()
+				if t.doneRotatingCh != nil {
+					t.doneRotatingCh <- struct{}{}
+				}
 			case <-t.closeCh:
 				return
 			}
diff --git a/lib/timedbs/timedbs_test.go b/lib/timedbs/timedbs_test.go
index 2126b9287..db21a5d79 100644
--- a/lib/timedbs/timedbs_test.go
+++ b/lib/timedbs/timedbs_test.go
@@ -1,21 +1,27 @@
-package timedbs_test
+package timedbs
 
 import (
 	"context"
 	"testing"
 	"time"
 
+	"github.com/raulk/clock"
 	"github.com/stretchr/testify/require"
 
 	blocks "github.com/ipfs/go-block-format"
 	"github.com/ipfs/go-cid"
-
-	"github.com/filecoin-project/lotus/lib/timedbs"
 )
 
 func TestTimedBSSimple(t *testing.T) {
-	tc := timedbs.NewTimedCacheBS(10 * time.Millisecond)
-	_ = tc.Start(context.Background())
+	tc := NewTimedCacheBS(10 * time.Millisecond)
+	mClock := clock.NewMock()
+	mClock.Set(time.Now())
+	tc.clock = mClock
+	tc.doneRotatingCh = make(chan struct{})
+
+	tc.Start(context.Background())
+	mClock.Add(1) // IDK why it is needed but it makes it work
+
 	defer func() {
 		_ = tc.Stop(context.Background())
 	}()
@@ -36,7 +42,8 @@ func TestTimedBSSimple(t *testing.T) {
 	require.NoError(t, err)
 	require.True(t, has)
 
-	time.Sleep(15 * time.Millisecond)
+	mClock.Add(10 * time.Millisecond)
+	<-tc.doneRotatingCh
 
 	// We should still have everything.
 	has, err = tc.Has(b1.Cid())
@@ -60,8 +67,8 @@ func TestTimedBSSimple(t *testing.T) {
 	require.NoError(t, err)
 	require.ElementsMatch(t, ks, []cid.Cid{b1.Cid(), b2.Cid(), b3.Cid()})
 
-	time.Sleep(10 * time.Millisecond)
-
+	mClock.Add(10 * time.Millisecond)
+	<-tc.doneRotatingCh
 	// should still have b2, and b3, but not b1
 
 	has, err = tc.Has(b1.Cid())