@@ -442,6 +442,8 @@ pub fn mock_env() -> Env {
442442/// ```
443443pub struct Envs {
444444 contract_address : Addr ,
445+ /// The number of nanoseconds between two consecutive blocks
446+ block_time : u64 ,
445447 last_height : u64 ,
446448 last_time : Timestamp ,
447449 envs_produced : u64 ,
@@ -455,21 +457,26 @@ impl Envs {
455457 Envs {
456458 // Default values here for compatibility with old `mock_env` function. They could be changed to anything else if there is a good reason.
457459 contract_address : api. addr_make ( "cosmos2contract" ) ,
460+ block_time : 5_000_000_000 , // 5s
458461 last_height : 12_344 ,
459462 last_time : Timestamp :: from_nanos ( 1_571_797_419_879_305_533 ) . minus_seconds ( 5 ) ,
460463 envs_produced : 0 ,
461464 }
462465 }
463466
464467 pub fn make ( & mut self ) -> Env {
465- let height = self . last_height + 1 ;
466- let time = self . last_time . plus_seconds ( 5 ) ;
468+ self . checked_make ( ) . unwrap ( )
469+ }
470+
471+ fn checked_make ( & mut self ) -> Option < Env > {
472+ let height = self . last_height . checked_add ( 1 ) ?;
473+ let time = Timestamp :: from_nanos ( self . last_time . nanos ( ) . checked_add ( self . block_time ) ?) ;
467474
468475 self . last_height = height;
469476 self . last_time = time;
470- self . envs_produced += 1 ;
477+ self . envs_produced += 1 ; // does not overflow because height increment fails first
471478
472- Env {
479+ Some ( Env {
473480 block : BlockInfo {
474481 height,
475482 time,
@@ -479,21 +486,17 @@ impl Envs {
479486 contract : ContractInfo {
480487 address : self . contract_address . clone ( ) ,
481488 } ,
482- }
489+ } )
483490 }
484491}
485492
486- // The iterator implementation can produce 1 million envs and then stops for no good reason.
493+ // The iterator implementation ends in case of overflows to avoid panics.
494+ // Using this is recommended for very long running test suites.
487495impl Iterator for Envs {
488496 type Item = Env ;
489497
490498 fn next ( & mut self ) -> Option < Self :: Item > {
491- if self . envs_produced < 1_000_000 {
492- let item = self . make ( ) ;
493- Some ( item)
494- } else {
495- None
496- }
499+ self . checked_make ( )
497500 }
498501}
499502
@@ -1427,7 +1430,7 @@ mod tests {
14271430 }
14281431
14291432 #[ test]
1430- fn envs_implements_iteratorworks ( ) {
1433+ fn envs_implements_iterator ( ) {
14311434 let envs = Envs :: new ( "food" ) ;
14321435
14331436 let result: Vec < _ > = envs. into_iter ( ) . take ( 5 ) . collect ( ) ;
@@ -1451,6 +1454,17 @@ mod tests {
14511454 result[ 4 ] . block. time,
14521455 Timestamp :: from_nanos( 1_571_797_439_879_305_533 )
14531456 ) ;
1457+
1458+ // Get a millions envs through iterator
1459+ let mut envs = Envs :: new ( "yo" ) ;
1460+ let first = envs. next ( ) . unwrap ( ) ;
1461+ let last = envs. take ( 1_000_000 ) . last ( ) . unwrap ( ) ;
1462+ assert_eq ! ( first. block. height, 12_345 ) ;
1463+ assert_eq ! ( last. block. height, 1_012_345 ) ;
1464+ assert_eq ! (
1465+ last. block. time,
1466+ first. block. time. plus_seconds( 1_000_000 * 5 )
1467+ ) ;
14541468 }
14551469
14561470 #[ test]
0 commit comments