github.com/kubeshop/testkube@v1.17.23/pkg/agent/events_test.go (about)

     1  package agent_test
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"net"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/kubeshop/testkube/pkg/executor/output"
    11  	"github.com/kubeshop/testkube/pkg/log"
    12  	"github.com/kubeshop/testkube/pkg/ui"
    13  
    14  	"github.com/stretchr/testify/assert"
    15  	"go.uber.org/zap"
    16  	"golang.org/x/sync/errgroup"
    17  	"google.golang.org/grpc"
    18  	"google.golang.org/grpc/metadata"
    19  
    20  	"github.com/kubeshop/testkube/internal/config"
    21  	"github.com/kubeshop/testkube/pkg/agent"
    22  	"github.com/kubeshop/testkube/pkg/api/v1/testkube"
    23  	"github.com/kubeshop/testkube/pkg/cloud"
    24  	"github.com/kubeshop/testkube/pkg/featureflags"
    25  )
    26  
    27  func TestEventLoop(t *testing.T) {
    28  	url := "localhost:8998"
    29  
    30  	ctx, cancel := context.WithCancel(context.Background())
    31  	cloudSrv := newEventServer(ctx)
    32  
    33  	go func() {
    34  		lis, err := net.Listen("tcp", url)
    35  		if err != nil {
    36  			panic(err)
    37  		}
    38  
    39  		var opts []grpc.ServerOption
    40  		grpcServer := grpc.NewServer(opts...)
    41  		cloud.RegisterTestKubeCloudAPIServer(grpcServer, cloudSrv)
    42  		err = grpcServer.Serve(lis)
    43  		if err != nil {
    44  			panic(err)
    45  		}
    46  	}()
    47  
    48  	logger, _ := zap.NewDevelopment()
    49  
    50  	grpcConn, err := agent.NewGRPCConnection(context.Background(), true, false, url, "", "", "", log.DefaultLogger)
    51  	ui.ExitOnError("error creating gRPC connection", err)
    52  	defer grpcConn.Close()
    53  
    54  	grpcClient := cloud.NewTestKubeCloudAPIClient(grpcConn)
    55  
    56  	var logStreamFunc func(ctx context.Context, executionID string) (chan output.Output, error)
    57  	var workflowNotificationsStreamFunc func(ctx context.Context, executionID string) (chan testkube.TestWorkflowExecutionNotification, error)
    58  
    59  	proContext := config.ProContext{APIKey: "api-key", WorkerCount: 5, LogStreamWorkerCount: 5, WorkflowNotificationsWorkerCount: 5}
    60  	agent, err := agent.NewAgent(logger.Sugar(), nil, grpcClient, logStreamFunc, workflowNotificationsStreamFunc, "", "", nil, featureflags.FeatureFlags{}, proContext)
    61  	assert.NoError(t, err)
    62  	go func() {
    63  		l, err := agent.Load()
    64  		if err != nil {
    65  			panic(err)
    66  		}
    67  
    68  		var i int
    69  		for {
    70  			res := l[0].Notify(testkube.Event{Id: fmt.Sprintf("%d", i)})
    71  			if res.Error_ != "" {
    72  				continue
    73  			}
    74  		}
    75  	}()
    76  
    77  	g, groupCtx := errgroup.WithContext(ctx)
    78  	g.Go(func() error {
    79  		return agent.Run(groupCtx)
    80  	})
    81  
    82  	time.Sleep(100 * time.Millisecond)
    83  	cancel()
    84  
    85  	g.Wait()
    86  
    87  	assert.True(t, cloudSrv.Count() >= 5)
    88  }
    89  
    90  func (cws *CloudEventServer) Count() int {
    91  	return cws.messageCount
    92  }
    93  
    94  func (cws *CloudEventServer) ExecuteAsync(srv cloud.TestKubeCloudAPI_ExecuteAsyncServer) error {
    95  	<-cws.ctx.Done()
    96  
    97  	return nil
    98  }
    99  
   100  func (cws *CloudEventServer) GetLogsStream(srv cloud.TestKubeCloudAPI_GetLogsStreamServer) error {
   101  	<-cws.ctx.Done()
   102  
   103  	return nil
   104  }
   105  
   106  func (cws *CloudEventServer) GetTestWorkflowNotificationsStream(srv cloud.TestKubeCloudAPI_GetTestWorkflowNotificationsStreamServer) error {
   107  	<-cws.ctx.Done()
   108  
   109  	return nil
   110  }
   111  
   112  func (cws *CloudEventServer) Send(srv cloud.TestKubeCloudAPI_SendServer) error {
   113  	md, ok := metadata.FromIncomingContext(srv.Context())
   114  	if !ok {
   115  		panic("no metadata")
   116  	}
   117  	apiKey := md.Get("api-key")
   118  	if apiKey[0] != "api-key" {
   119  		panic("error bad api-key")
   120  	}
   121  
   122  	for {
   123  		if srv.Context().Err() != nil {
   124  			return srv.Context().Err()
   125  		}
   126  		resp, err := srv.Recv()
   127  		if err != nil {
   128  			return err
   129  		}
   130  
   131  		if resp.Opcode == cloud.Opcode_HEALTH_CHECK {
   132  			continue
   133  		}
   134  
   135  		if resp.Opcode != cloud.Opcode_TEXT_FRAME {
   136  			panic("bad opcode")
   137  		}
   138  		cws.messageCount++
   139  
   140  		if cws.messageCount >= 5 {
   141  			return nil
   142  		}
   143  	}
   144  }
   145  
   146  func newEventServer(ctx context.Context) *CloudEventServer {
   147  	return &CloudEventServer{ctx: ctx}
   148  }
   149  
   150  type CloudEventServer struct {
   151  	cloud.UnimplementedTestKubeCloudAPIServer
   152  	messageCount int
   153  	ctx          context.Context
   154  }