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