From 4cb81a0f5d19763514c28a191e5e8f7083a6b890 Mon Sep 17 00:00:00 2001
From: Naohiro Yoshida <naohiro.yoshida@datachain.jp>
Date: Thu, 2 Nov 2023 09:25:59 +0900
Subject: [PATCH] add test

Signed-off-by: Naohiro Yoshida <naohiro.yoshida@datachain.jp>
---
 module/prover_test.go | 63 +++++++++++++++++++++++++++++++++++++++----
 1 file changed, 58 insertions(+), 5 deletions(-)

diff --git a/module/prover_test.go b/module/prover_test.go
index a39c4a3..fafe381 100644
--- a/module/prover_test.go
+++ b/module/prover_test.go
@@ -35,8 +35,10 @@ const (
 
 type mockChain struct {
 	Chain
-	latestHeight uint64
-	chainID      uint64
+	latestHeight            uint64
+	consensusStateTimestamp map[exported.Height]uint64
+	chainTimestamp          map[exported.Height]uint64
+	chainID                 uint64
 }
 
 func (c *mockChain) CanonicalChainID(ctx context.Context) (uint64, error) {
@@ -56,6 +58,21 @@ func (c *mockChain) QueryClientState(ctx core.QueryContext) (*clienttypes.QueryC
 	return clienttypes.NewQueryClientStateResponse(anyClientState, nil, cHeight), nil
 }
 
+func (c *mockChain) QueryClientConsensusState(ctx core.QueryContext, height exported.Height) (*clienttypes.QueryConsensusStateResponse, error) {
+	cHeight := clienttypes.NewHeight(ctx.Height().GetRevisionNumber(), ctx.Height().GetRevisionHeight())
+	cs := ConsensusState{
+		StateRoot:              common.Hash{}.Bytes(),
+		Timestamp:              c.consensusStateTimestamp[height],
+		CurrentValidatorsHash:  common.Hash{}.Bytes(),
+		PreviousValidatorsHash: common.Hash{}.Bytes(),
+	}
+	anyConsensusState, err := codectypes.NewAnyWithValue(&cs)
+	if err != nil {
+		return nil, err
+	}
+	return clienttypes.NewQueryConsensusStateResponse(anyConsensusState, nil, cHeight), nil
+}
+
 func (c *mockChain) Header(_ context.Context, height uint64) (*types2.Header, error) {
 	headerMap := map[uint64]types2.Header{
 		31297221: {
@@ -258,6 +275,10 @@ func (c *mockChain) LatestHeight() (exported.Height, error) {
 	return clienttypes.NewHeight(0, c.latestHeight), nil
 }
 
+func (c *mockChain) Timestamp(height exported.Height) (time.Time, error) {
+	return time.Unix(int64(c.chainTimestamp[height]), 0), nil
+}
+
 type ProverTestSuite struct {
 	suite.Suite
 	prover *Prover
@@ -296,11 +317,17 @@ func (ts *ProverTestSuite) SetupTest() {
 	config := ProverConfig{
 		TrustingPeriod: 100 * time.Second,
 		MaxClockDrift:  1 * time.Millisecond,
+		RefreshThresholdRate: &Fraction{
+			Numerator:   1,
+			Denominator: 2,
+		},
 	}
 	ts.chain = &mockChain{
-		Chain:        NewChain(chain),
-		latestHeight: 31297221,
-		chainID:      9999,
+		Chain:                   NewChain(chain),
+		latestHeight:            31297221,
+		chainID:                 9999,
+		consensusStateTimestamp: make(map[exported.Height]uint64),
+		chainTimestamp:          make(map[exported.Height]uint64),
 	}
 	ts.prover = NewProver(ts.chain, &config).(*Prover)
 }
@@ -462,3 +489,29 @@ func (ts *ProverTestSuite) TestConnectionStateProofAsLCPCommitment() {
 	ts.Require().Equal(commitmentHeight, uint64(317))
 	ts.Require().Equal(commitmentStateId.String(), "0xee0b5f32ae2bff0d82149ea22b02e350fbbe467a514ba80bbadd89007df1d167")
 }
+
+func (ts *ProverTestSuite) TestCheckRefreshRequired() {
+	type dstMock struct {
+		Chain
+		core.Prover
+	}
+	dst := dstMock{
+		Chain:  ts.prover.chain,
+		Prover: ts.prover,
+	}
+	now := time.Now()
+	height := clienttypes.NewHeight(0, ts.chain.latestHeight)
+	ts.chain.chainTimestamp[height] = uint64(now.Unix())
+
+	// should refresh
+	ts.chain.consensusStateTimestamp[height] = uint64(now.Add(-51 * time.Second).Unix())
+	required, err := ts.prover.CheckRefreshRequired(dst)
+	ts.Require().NoError(err)
+	ts.Require().True(required)
+
+	// needless
+	ts.chain.consensusStateTimestamp[height] = uint64(now.Add(-50 * time.Second).Unix())
+	required, err = ts.prover.CheckRefreshRequired(dst)
+	ts.Require().NoError(err)
+	ts.Require().False(required)
+}