@@ -557,15 +557,15 @@ func TestRequestCollation_Timeout(t *testing.T) {
557557 start := time .Now ()
558558
559559 // This should timeout after ~1 second
560- collation , err := cpvs .requestCollation (relayParent , paraID , peerID )
560+ collation , err := cpvs .requestCollation (relayParent , paraID , peerID , nil )
561561
562562 elapsed := time .Since (start )
563563
564564 // Verify timeout behavior
565565 require .Nil (t , collation )
566566 require .Error (t , err )
567567 require .Contains (t , err .Error (), "timed out" )
568- require .Greater (t , elapsed , 900 * time .Millisecond )
568+ require .Greater (t , elapsed , 90 * time .Millisecond )
569569 require .Less (t , elapsed , 1200 * time .Millisecond )
570570}
571571
@@ -605,7 +605,7 @@ func TestRequestCollation_Success(t *testing.T) {
605605 peerID := peer .ID ("test-peer" )
606606
607607 start := time .Now ()
608- _ , err := cpvs .requestCollation (relayParent , paraID , peerID )
608+ _ , err := cpvs .requestCollation (relayParent , paraID , peerID , nil )
609609 elapsed := time .Since (start )
610610
611611 // Test that it completed quickly (didn't timeout)
@@ -626,8 +626,57 @@ func TestRequestCollation_OutOfView(t *testing.T) {
626626 paraID := parachaintypes .ParaID (123 )
627627 peerID := peer .ID ("test-peer" )
628628
629- collation , err := cpvs .requestCollation (relayParent , paraID , peerID )
629+ collation , err := cpvs .requestCollation (relayParent , paraID , peerID , nil )
630630
631631 require .Nil (t , collation )
632632 require .Equal (t , ErrOutOfView , err )
633633}
634+
635+ func TestPeerVersionManagement (t * testing.T ) {
636+ t .Parallel ()
637+
638+ cpvs := & CollatorProtocolValidatorSide {
639+ peerVersions : make (map [peer.ID ]PeerProtocolVersion ),
640+ }
641+
642+ peerID := peer .ID ("test-peer" )
643+
644+ // Test default version
645+ version := cpvs .getPeerProtocolVersion (peerID )
646+ require .Equal (t , ProtocolV1 , version , "New peer should default to V1" )
647+
648+ // Test setting V2
649+ cpvs .setPeerProtocolVersion (peerID , ProtocolV2 )
650+ version = cpvs .getPeerProtocolVersion (peerID )
651+ require .Equal (t , ProtocolV2 , version , "Peer should be upgraded to V2" )
652+ }
653+
654+ func TestV2PeerRequiresCandidateHash (t * testing.T ) {
655+ t .Parallel ()
656+
657+ ctrl := gomock .NewController (t )
658+ defer ctrl .Finish ()
659+
660+ mockNet := NewMockNetwork (ctrl )
661+ mockNet .EXPECT ().
662+ GetRequestResponseProtocol (gomock .Any (), gomock .Any (), gomock .Any ()).
663+ Return (& network.RequestResponseProtocol {}).
664+ AnyTimes ()
665+
666+ cpvs := New (mockNet , protocol .ID ("/test/collations/1" ), make (chan any , 10 ), nil , nil )
667+
668+ // Setup
669+ relayParent := common.Hash {0x01 }
670+ cpvs .perRelayParent = map [common.Hash ]PerRelayParent {relayParent : {}}
671+ cpvs .collationRequests = make (chan CollationRequestInfo , 100 )
672+
673+ peerID := peer .ID ("v2-peer" )
674+
675+ // Set peer to V2
676+ cpvs .setPeerProtocolVersion (peerID , ProtocolV2 )
677+
678+ // Test: V2 peer without candidate hash should fail
679+ _ , err := cpvs .requestCollation (relayParent , 123 , peerID , nil )
680+ require .Error (t , err )
681+ require .Contains (t , err .Error (), "V2 peer requires candidate hash" )
682+ }
0 commit comments