77 "fmt"
88 "os"
99 "strconv"
10+ "sync"
11+ "sync/atomic"
1012 "testing"
1113 "time"
1214
@@ -160,6 +162,7 @@ func (m *MockBrokerChainBridge) GetParams() (*types.QueryParamsResponse, error)
160162
161163type MockRandomSeedManager struct {
162164 mock.Mock
165+ onGenerate func (epochIndex uint64 )
163166}
164167
165168func (m * MockRandomSeedManager ) ChangeCurrentSeed () {
@@ -176,16 +179,28 @@ func (m *MockRandomSeedManager) RequestMoney(epochIndex uint64) {
176179}
177180
178181func (m * MockRandomSeedManager ) CreateNewSeed (epochIndex uint64 ) (* apiconfig.SeedInfo , error ) {
179- m .Called ()
180- return nil , nil
182+ m .Called (epochIndex )
183+ // Match the signature ListRandomSeeds returns after GenerateSeedInfo so
184+ // confirmSeedLocally can succeed the same way production restore does.
185+ return & apiconfig.SeedInfo {
186+ Seed : 1 ,
187+ EpochIndex : epochIndex ,
188+ Signature : integrationTestSeedSignature ,
189+ }, nil
181190}
182191
183192func (m * MockRandomSeedManager ) GenerateSeedInfo (epochIndex uint64 ) {
184193 m .Called (epochIndex )
194+ if m .onGenerate != nil {
195+ m .onGenerate (epochIndex )
196+ }
185197}
186198
187199type MockQueryClient struct {
188200 mock.Mock
201+ listRandomSeedsCalls atomic.Int64
202+ mu sync.Mutex
203+ submittedByEpoch map [uint64 ]* types.RandomSeed
189204}
190205
191206func (m * MockQueryClient ) EpochInfo (ctx context.Context , req * types.QueryEpochInfoRequest , opts ... grpc.CallOption ) (* types.QueryEpochInfoResponse , error ) {
@@ -201,6 +216,36 @@ func (m *MockQueryClient) Params(ctx context.Context, req *types.QueryParamsRequ
201216 return args .Get (0 ).(* types.QueryParamsResponse ), args .Error (1 )
202217}
203218
219+ const integrationTestSeedParticipant = "some-address"
220+ const integrationTestSeedSignature = "integration-test-seed-signature"
221+
222+ // ListRandomSeeds is not testify-expectation based: ensureSeedSubmitted runs
223+ // asynchronously and can race ExpectedCalls = nil in setLatestEpoch.
224+ // After GenerateSeedInfo, the seed becomes visible here so later ensures take
225+ // the confirm path and stop resubmitting (same shape as production).
226+ func (m * MockQueryClient ) ListRandomSeeds (ctx context.Context , req * types.QueryRandomSeedsRequest , opts ... grpc.CallOption ) (* types.QueryRandomSeedsResponse , error ) {
227+ m .listRandomSeedsCalls .Add (1 )
228+ m .mu .Lock ()
229+ defer m .mu .Unlock ()
230+ if seed , ok := m .submittedByEpoch [req .EpochIndex ]; ok {
231+ return & types.QueryRandomSeedsResponse {Seeds : []* types.RandomSeed {seed }}, nil
232+ }
233+ return & types.QueryRandomSeedsResponse {}, nil
234+ }
235+
236+ func (m * MockQueryClient ) markSeedSubmitted (epochIndex uint64 ) {
237+ m .mu .Lock ()
238+ defer m .mu .Unlock ()
239+ if m .submittedByEpoch == nil {
240+ m .submittedByEpoch = make (map [uint64 ]* types.RandomSeed )
241+ }
242+ m .submittedByEpoch [epochIndex ] = & types.RandomSeed {
243+ Participant : integrationTestSeedParticipant ,
244+ EpochIndex : epochIndex ,
245+ Signature : integrationTestSeedSignature ,
246+ }
247+ }
248+
204249// Test setup helpers
205250
206251type IntegrationTestSetup struct {
@@ -219,7 +264,11 @@ func createIntegrationTestSetup(reconcilialtionConfig *MlNodeReconciliationConfi
219264 os .Setenv ("ENFORCED_MODEL_ID" , "disabled" )
220265
221266 mockQueryClient := & MockQueryClient {}
222- mockSeedManager := & MockRandomSeedManager {}
267+ mockSeedManager := & MockRandomSeedManager {
268+ onGenerate : func (epochIndex uint64 ) {
269+ mockQueryClient .markSeedSubmitted (epochIndex )
270+ },
271+ }
223272
224273 phaseTracker := & chainphase.ChainPhaseTracker {}
225274
@@ -446,7 +495,19 @@ func (setup *IntegrationTestSetup) simulateBlock(height int64) error {
446495 Height : height ,
447496 Hash : fmt .Sprintf ("hash-%d" , height ),
448497 }
449- return setup .Dispatcher .ProcessNewBlock (context .Background (), blockInfo )
498+ err := setup .Dispatcher .ProcessNewBlock (context .Background (), blockInfo )
499+ setup .waitForSeedEnsureIdle ()
500+ return err
501+ }
502+
503+ func (setup * IntegrationTestSetup ) waitForSeedEnsureIdle () {
504+ deadline := time .Now ().Add (2 * time .Second )
505+ for time .Now ().Before (deadline ) {
506+ if ! setup .Dispatcher .seedEnsureInFlight .Load () {
507+ return
508+ }
509+ time .Sleep (5 * time .Millisecond )
510+ }
450511}
451512
452513func (setup * IntegrationTestSetup ) getNodeClient (nodeId string , port int ) * mlnodeclient.MockClient {
@@ -653,6 +714,8 @@ func TestRegularPocScenario(t *testing.T) {
653714 assertNodeClient (t , expected , node2Client )
654715 i ++
655716 }
717+ require .GreaterOrEqual (t , setup .MockQueryClient .listRandomSeedsCalls .Load (), int64 (1 ),
718+ "seed ensure should query ListRandomSeeds at least once during PoC window" )
656719
657720 pocValStart := i
658721 pocValEnd := pocValStart + setup .EpochParams .PocValidationDelay + setup .EpochParams .PocValidationDuration
0 commit comments