github.com/gravitational/teleport/api@v0.0.0-20240507183017-3110591cbafc/client/inventory_test.go (about)

     1  /*
     2  Copyright 2022 Gravitational, Inc.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package client
    18  
    19  import (
    20  	"context"
    21  	"io"
    22  	"testing"
    23  	"time"
    24  
    25  	"github.com/stretchr/testify/assert"
    26  	"github.com/stretchr/testify/require"
    27  
    28  	"github.com/gravitational/teleport/api/client/proto"
    29  )
    30  
    31  // TestInventoryControlStreamPipe is a sanity-check to make sure that the in-memory
    32  // pipe version of the ICS works as expected.  This test is trivial but it helps to
    33  // keep accidental breakage of the pipe abstraction from showing up in an obscure
    34  // way inside the tests that rely upon it.
    35  func TestInventoryControlStreamPipe(t *testing.T) {
    36  	ctx, cancel := context.WithCancel(context.Background())
    37  	defer cancel()
    38  
    39  	upstream, downstream := InventoryControlStreamPipe()
    40  	defer upstream.Close()
    41  
    42  	upMsgs := []proto.UpstreamInventoryMessage{
    43  		proto.UpstreamInventoryHello{},
    44  		proto.UpstreamInventoryPong{},
    45  		proto.InventoryHeartbeat{},
    46  	}
    47  
    48  	downMsgs := []proto.DownstreamInventoryMessage{
    49  		proto.DownstreamInventoryHello{},
    50  		proto.DownstreamInventoryPing{},
    51  		proto.DownstreamInventoryPing{}, // duplicate to pad downMsgs to same length as upMsgs
    52  	}
    53  
    54  	go func() {
    55  		for _, m := range upMsgs {
    56  			downstream.Send(ctx, m)
    57  		}
    58  	}()
    59  
    60  	go func() {
    61  		for _, m := range downMsgs {
    62  			upstream.Send(ctx, m)
    63  		}
    64  	}()
    65  
    66  	timeout := time.NewTimer(time.Second * 5)
    67  	defer timeout.Stop()
    68  	for i := range upMsgs {
    69  		if !timeout.Stop() {
    70  			<-timeout.C
    71  		}
    72  		timeout.Reset(time.Second * 5)
    73  
    74  		// upstream handle recv
    75  		select {
    76  		case msg := <-upstream.Recv():
    77  			require.IsType(t, upMsgs[i], msg)
    78  		case <-timeout.C:
    79  			t.Fatalf("timeout waiting for message: %T", upMsgs[i])
    80  		}
    81  
    82  		// downstream handle recv
    83  		select {
    84  		case msg := <-downstream.Recv():
    85  			require.IsType(t, downMsgs[i], msg)
    86  		case <-timeout.C:
    87  			t.Fatalf("timeout waiting for message: %T", downMsgs[i])
    88  		}
    89  	}
    90  
    91  	upstream.Close()
    92  
    93  	if !timeout.Stop() {
    94  		<-timeout.C
    95  	}
    96  	timeout.Reset(time.Second * 5)
    97  
    98  	select {
    99  	case <-downstream.Done():
   100  	case <-timeout.C:
   101  		t.Fatal("timeout waiting for close")
   102  	}
   103  
   104  	assert.ErrorIs(t, downstream.Error(), io.EOF)
   105  }