github.com/cilium/cilium@v1.16.2/pkg/envoy/xds/stream_test.go (about)

     1  // SPDX-License-Identifier: Apache-2.0
     2  // Copyright Authors of Cilium
     3  package xds
     4  
     5  import (
     6  	"context"
     7  	"errors"
     8  	"io"
     9  	"time"
    10  
    11  	envoy_service_discovery "github.com/cilium/proxy/go/envoy/service/discovery/v3"
    12  )
    13  
    14  // MockStream is a mock implementation of Stream used for testing.
    15  type MockStream struct {
    16  	ctx         context.Context
    17  	recv        chan *envoy_service_discovery.DiscoveryRequest
    18  	sent        chan *envoy_service_discovery.DiscoveryResponse
    19  	recvTimeout time.Duration
    20  	sentTimeout time.Duration
    21  }
    22  
    23  // NewMockStream creates a new mock Stream for testing.
    24  func NewMockStream(ctx context.Context, recvSize, sentSize int, recvTimeout, sentTimeout time.Duration) *MockStream {
    25  	return &MockStream{
    26  		ctx:         ctx,
    27  		recv:        make(chan *envoy_service_discovery.DiscoveryRequest, recvSize),
    28  		sent:        make(chan *envoy_service_discovery.DiscoveryResponse, sentSize),
    29  		recvTimeout: recvTimeout,
    30  		sentTimeout: sentTimeout,
    31  	}
    32  }
    33  
    34  func (s *MockStream) Send(resp *envoy_service_discovery.DiscoveryResponse) error {
    35  	subCtx, cancel := context.WithTimeout(s.ctx, s.sentTimeout)
    36  
    37  	select {
    38  	case <-subCtx.Done():
    39  		cancel()
    40  		if errors.Is(subCtx.Err(), context.Canceled) {
    41  			return io.EOF
    42  		}
    43  		return subCtx.Err()
    44  	case s.sent <- resp:
    45  		cancel()
    46  		return nil
    47  	}
    48  }
    49  
    50  func (s *MockStream) Recv() (*envoy_service_discovery.DiscoveryRequest, error) {
    51  	subCtx, cancel := context.WithTimeout(s.ctx, s.recvTimeout)
    52  
    53  	select {
    54  	case <-subCtx.Done():
    55  		cancel()
    56  		if errors.Is(subCtx.Err(), context.Canceled) {
    57  			return nil, io.EOF
    58  		}
    59  		return nil, subCtx.Err()
    60  	case req := <-s.recv:
    61  		cancel()
    62  		if req == nil {
    63  			return nil, io.EOF
    64  		}
    65  		return req, nil
    66  	}
    67  }
    68  
    69  // SendRequest queues a request to be received by calling Recv.
    70  func (s *MockStream) SendRequest(req *envoy_service_discovery.DiscoveryRequest) error {
    71  	subCtx, cancel := context.WithTimeout(s.ctx, s.recvTimeout)
    72  
    73  	select {
    74  	case <-subCtx.Done():
    75  		cancel()
    76  		if errors.Is(subCtx.Err(), context.Canceled) {
    77  			return io.EOF
    78  		}
    79  		return subCtx.Err()
    80  	case s.recv <- req:
    81  		cancel()
    82  		return nil
    83  	}
    84  }
    85  
    86  // RecvResponse receives a response that was queued by calling Send.
    87  func (s *MockStream) RecvResponse() (*envoy_service_discovery.DiscoveryResponse, error) {
    88  	subCtx, cancel := context.WithTimeout(s.ctx, s.sentTimeout)
    89  
    90  	select {
    91  	case <-subCtx.Done():
    92  		cancel()
    93  		if errors.Is(subCtx.Err(), context.Canceled) {
    94  			return nil, io.EOF
    95  		}
    96  		return nil, subCtx.Err()
    97  	case resp := <-s.sent:
    98  		cancel()
    99  		return resp, nil
   100  	}
   101  }
   102  
   103  // Close closes the resources used by this MockStream.
   104  func (s *MockStream) Close() {
   105  	close(s.recv)
   106  	close(s.sent)
   107  }