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

     1  package agent_test
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"net"
     7  	"sync/atomic"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/kubeshop/testkube/pkg/api/v1/testkube"
    12  	"github.com/kubeshop/testkube/pkg/executor/output"
    13  	"github.com/kubeshop/testkube/pkg/log"
    14  	"github.com/kubeshop/testkube/pkg/ui"
    15  
    16  	"github.com/stretchr/testify/assert"
    17  	"github.com/valyala/fasthttp"
    18  	"go.uber.org/zap"
    19  	"golang.org/x/sync/errgroup"
    20  	"google.golang.org/grpc"
    21  	"google.golang.org/grpc/metadata"
    22  
    23  	"github.com/kubeshop/testkube/internal/config"
    24  	"github.com/kubeshop/testkube/pkg/agent"
    25  	"github.com/kubeshop/testkube/pkg/cloud"
    26  	"github.com/kubeshop/testkube/pkg/featureflags"
    27  )
    28  
    29  func TestCommandExecution(t *testing.T) {
    30  	url := "localhost:9999"
    31  	ctx, cancel := context.WithCancel(context.Background())
    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, newServer(ctx))
    42  		grpcServer.Serve(lis)
    43  	}()
    44  
    45  	var msgCnt int32
    46  	m := func(ctx *fasthttp.RequestCtx) {
    47  		h := &ctx.Response.Header
    48  		h.Add("Content-type", "application/json")
    49  		fmt.Fprintf(ctx, "Hi there! RequestURI is %q", ctx.RequestURI())
    50  		atomic.AddInt32(&msgCnt, 1)
    51  	}
    52  
    53  	grpcConn, err := agent.NewGRPCConnection(context.Background(), true, false, url, "", "", "", log.DefaultLogger)
    54  	ui.ExitOnError("error creating gRPC connection", err)
    55  	defer grpcConn.Close()
    56  
    57  	grpcClient := cloud.NewTestKubeCloudAPIClient(grpcConn)
    58  
    59  	var logStreamFunc func(ctx context.Context, executionID string) (chan output.Output, error)
    60  	var workflowNotificationsStreamFunc func(ctx context.Context, executionID string) (chan testkube.TestWorkflowExecutionNotification, error)
    61  
    62  	logger, _ := zap.NewDevelopment()
    63  	proContext := config.ProContext{APIKey: "api-key", WorkerCount: 5, LogStreamWorkerCount: 5, WorkflowNotificationsWorkerCount: 5}
    64  	agent, err := agent.NewAgent(logger.Sugar(), m, grpcClient, logStreamFunc, workflowNotificationsStreamFunc, "", "", nil, featureflags.FeatureFlags{}, proContext)
    65  	if err != nil {
    66  		t.Fatal(err)
    67  	}
    68  
    69  	g, groupCtx := errgroup.WithContext(ctx)
    70  	g.Go(func() error {
    71  		return agent.Run(groupCtx)
    72  	})
    73  
    74  	time.Sleep(100 * time.Millisecond)
    75  	cancel()
    76  
    77  	g.Wait()
    78  
    79  	assert.True(t, msgCnt > 0)
    80  }
    81  
    82  type CloudServer struct {
    83  	cloud.UnimplementedTestKubeCloudAPIServer
    84  	ctx context.Context
    85  }
    86  
    87  func (cs *CloudServer) GetLogsStream(srv cloud.TestKubeCloudAPI_GetLogsStreamServer) error {
    88  	<-cs.ctx.Done()
    89  
    90  	return nil
    91  }
    92  
    93  func (cs *CloudServer) GetTestWorkflowNotificationsStream(srv cloud.TestKubeCloudAPI_GetTestWorkflowNotificationsStreamServer) error {
    94  	<-cs.ctx.Done()
    95  
    96  	return nil
    97  }
    98  
    99  func (cs *CloudServer) ExecuteAsync(srv cloud.TestKubeCloudAPI_ExecuteAsyncServer) error {
   100  	md, ok := metadata.FromIncomingContext(srv.Context())
   101  	if !ok {
   102  		panic("no metadata")
   103  	}
   104  	apiKey := md.Get("api-key")
   105  	if apiKey[0] != "api-key" {
   106  		panic("error bad api-key")
   107  	}
   108  
   109  	req := &cloud.ExecuteRequest{Method: "url", Url: "localhost/v1/tests", Body: nil}
   110  	err := srv.Send(req)
   111  	if err != nil {
   112  		return err
   113  	}
   114  
   115  	resp, err := srv.Recv()
   116  	if err != nil {
   117  		return err
   118  	}
   119  	fmt.Println(resp)
   120  
   121  	return nil
   122  }
   123  
   124  func (cs *CloudServer) Send(srv cloud.TestKubeCloudAPI_SendServer) error {
   125  	for {
   126  		if srv.Context().Err() != nil {
   127  			return srv.Context().Err()
   128  		}
   129  
   130  		_, err := srv.Recv()
   131  		if err != nil {
   132  			return err
   133  		}
   134  	}
   135  
   136  }
   137  func newServer(ctx context.Context) *CloudServer {
   138  	return &CloudServer{ctx: ctx}
   139  }