@@ -29,6 +29,7 @@ import (
2929
3030 "github.com/api7/gopkg/pkg/log"
3131 "github.com/gavv/httpexpect/v2"
32+ "github.com/gorilla/websocket"
3233 "github.com/gruntwork-io/terratest/modules/k8s"
3334 "github.com/gruntwork-io/terratest/modules/testing"
3435 . "github.com/onsi/ginkgo/v2" //nolint:staticcheck
@@ -579,3 +580,141 @@ func (s *Scaffold) GetMetricsEndpoint() string {
579580 s .addFinalizers (tunnel .Close )
580581 return fmt .Sprintf ("http://%s/metrics" , tunnel .Endpoint ())
581582}
583+ << << << < HEAD
584+ == == == =
585+
586+ func (s * Scaffold ) NewWebsocketClient (tls * tls.Config , path string , headers http.Header ) * websocket.Conn {
587+ var host = s .ApisixHTTPEndpoint ()
588+ var scheme = "ws"
589+ if tls != nil {
590+ scheme = "wss"
591+ host = s .GetAPISIXHTTPSEndpoint ()
592+ }
593+
594+ dialer := websocket.Dialer {
595+ TLSClientConfig : tls ,
596+ }
597+
598+ u := url.URL {
599+ Scheme : scheme ,
600+ Host : host ,
601+ Path : path ,
602+ }
603+ var conn * websocket.Conn
604+
605+ s .RetryAssertion (func () error {
606+ c , resp , err := dialer .Dial (u .String (), headers )
607+ if err != nil {
608+ return err
609+ }
610+ if resp == nil || resp .StatusCode != http .StatusSwitchingProtocols {
611+ _ = c .Close ()
612+ return fmt .Errorf ("unexpected status code: %d" , resp .StatusCode )
613+ }
614+ conn = c
615+ return nil
616+ }).ShouldNot (HaveOccurred (), "establishing websocket connection" )
617+ return conn
618+ }
619+
620+ func (s * Scaffold ) ControlAPIClient () (ControlAPIClient , error ) {
621+ tunnel := k8s .NewTunnel (s .kubectlOptions , k8s .ResourceTypeService , "apisix-control-api" , 9090 , 9090 )
622+ if err := tunnel .ForwardPortE (s .t ); err != nil {
623+ return nil , err
624+ }
625+ s .addFinalizers (tunnel .Close )
626+
627+ return & controlAPI {
628+ client : NewClient ("http" , tunnel .Endpoint ()),
629+ }, nil
630+ }
631+
632+ func (s * Scaffold ) EnsureNumService (controlAPIClient ControlAPIClient , matcher func (result int ) bool ) error {
633+ times := 0
634+ return wait .PollUntilContextTimeout (context .Background (), 100 * time .Millisecond , 10 * time .Minute , true , func (ctx context.Context ) (done bool , err error ) {
635+ times ++
636+ results , _ , err := controlAPIClient .ListServices ()
637+ if err != nil {
638+ log .Errorw ("failed to ListServices" , zap .Error (err ))
639+ return false , nil
640+ }
641+ if ! matcher (len (results )) {
642+ log .Debugw ("number of effective services" , zap .Int ("number" , len (results )), zap .Int ("times" , times ))
643+ return false , nil
644+ }
645+ return true , nil
646+ })
647+ }
648+
649+ func (s * Scaffold ) ExpectUpstream (controlAPIClient ControlAPIClient , name string , matcher func (upstream adctypes.Upstream ) bool ) error {
650+ times := 0
651+ return wait .PollUntilContextTimeout (context .Background (), 1 * time .Second , 10 * time .Minute , true , func (ctx context.Context ) (done bool , err error ) {
652+ times ++
653+ upstreams , _ , err := controlAPIClient .ListUpstreams ()
654+ if err != nil {
655+ log .Errorw ("failed to ListServices" , zap .Error (err ))
656+ return false , nil
657+ }
658+ for _ , upstream := range upstreams {
659+ upsValue := upstream .(map [string ]any )
660+ data , err := json .Marshal (upsValue ["value" ])
661+ if err != nil {
662+ return false , fmt .Errorf ("failed to marshal upstream: %v" , err )
663+ }
664+
665+ var ups adctypes.Upstream
666+ if err := json .Unmarshal (data , & ups ); err != nil {
667+ return false , fmt .Errorf ("failed to unmarshal upstream: %v" , err )
668+ }
669+ if name != "" && ups .Name != name {
670+ continue
671+ }
672+ if ok := matcher (ups ); ! ok {
673+ return false , nil
674+ }
675+ }
676+ return true , nil
677+ })
678+ }
679+
680+ func (s * Scaffold ) EnsureNumUpstreamNodes (controlAPIClient ControlAPIClient , name string , number int ) error {
681+ return s .ExpectUpstream (controlAPIClient , name , func (upstream adctypes.Upstream ) bool {
682+ if len (upstream .Nodes ) != number {
683+ log .Warnf ("expect upstream: [%s] nodes num to be %d, but got %d" , upstream .Name , number , len (upstream .Nodes ))
684+ return false
685+ }
686+ return true
687+ })
688+ }
689+
690+ type ControlAPIClient interface {
691+ ListServices () ([]any , int64 , error )
692+ ListUpstreams () ([]any , int64 , error )
693+ }
694+
695+ type controlAPI struct {
696+ client * httpexpect.Expect
697+ }
698+
699+ func (c * controlAPI ) ListUpstreams () (result []any , total int64 , err error ) {
700+ resp := c .client .Request (http .MethodGet , "/v1/upstreams" ).Expect ()
701+ if resp .Raw ().StatusCode != http .StatusOK {
702+ return nil , 0 , fmt .Errorf ("unexpected status code: %v, message: %s" , resp .Raw ().StatusCode , resp .Body ().Raw ())
703+ }
704+ if err = json .Unmarshal ([]byte (resp .Body ().Raw ()), & result ); err != nil {
705+ return nil , 0 , fmt .Errorf ("failed to unmarshal response body: %w" , err )
706+ }
707+ return result , int64 (len (result )), err
708+ }
709+
710+ func (c * controlAPI ) ListServices () (result []any , total int64 , err error ) {
711+ resp := c .client .Request (http .MethodGet , "/v1/services" ).Expect ()
712+ if resp .Raw ().StatusCode != http .StatusOK {
713+ return nil , 0 , fmt .Errorf ("unexpected status code: %v, message: %s" , resp .Raw ().StatusCode , resp .Body ().Raw ())
714+ }
715+ if err = json .Unmarshal ([]byte (resp .Body ().Raw ()), & result ); err != nil {
716+ return nil , 0 , fmt .Errorf ("failed to unmarshal response body: %w" , err )
717+ }
718+ return result , int64 (len (result )), err
719+ }
720+ >> >> >> > b70e22eb (fix : wss related tests are unstable (#2675 ))
![backport-bot[bot]](https://avatars.githubusercontent.com/in/43169?v=4&size=40){kind=avatar}
0 commit comments