chore(server): Updating data store test to use queues (#3166)

* chore(server): Updating data store test to use queues

* chore(server): Updating data store test to use queues

* fix: build

* Adding postgres pipelines

* fix: agent unit tests

* using testconnection job instead of executor job

* cleanup

* PR comments

Author: xoscar

agent/client/client.go

@@ -26,9 +26,10 @@ type Client struct {
 	sessionConfig *SessionConfig
 	done          chan bool
 
-	triggerListener  func(context.Context, *proto.TriggerRequest) error
-	pollListener     func(context.Context, *proto.PollingRequest) error
-	shutdownListener func(context.Context, *proto.ShutdownRequest) error
+	triggerListener             func(context.Context, *proto.TriggerRequest) error
+	pollListener                func(context.Context, *proto.PollingRequest) error
+	shutdownListener            func(context.Context, *proto.ShutdownRequest) error
+	dataStoreConnectionListener func(context.Context, *proto.DataStoreConnectionTestRequest) error
 }
 
 func (c *Client) Start(ctx context.Context) error {
@@ -61,6 +62,11 @@ func (c *Client) Start(ctx context.Context) error {
 		return err
 	}
 
+	err = c.startDataStoreConnectionTestListener(ctx)
+	if err != nil {
+		return err
+	}
+
 	c.startHearthBeat(ctx)
 
 	return nil
@@ -88,6 +94,10 @@ func (c *Client) OnTriggerRequest(listener func(context.Context, *proto.TriggerR
 	c.triggerListener = listener
 }
 
+func (c *Client) OnDataStoreTestConnectionRequest(listener func(context.Context, *proto.DataStoreConnectionTestRequest) error) {
+	c.dataStoreConnectionListener = listener
+}
+
 func (c *Client) OnPollingRequest(listener func(context.Context, *proto.PollingRequest) error) {
 	c.pollListener = listener
 }

agent/client/mocks/grpc_server.go

@@ -13,20 +13,22 @@ import (
 
 type GrpcServerMock struct {
 	proto.UnimplementedOrchestratorServer
-	port               int
-	triggerChannel     chan *proto.TriggerRequest
-	pollingChannel     chan *proto.PollingRequest
-	terminationChannel chan *proto.ShutdownRequest
+	port                 int
+	triggerChannel       chan *proto.TriggerRequest
+	pollingChannel       chan *proto.PollingRequest
+	terminationChannel   chan *proto.ShutdownRequest
+	dataStoreTestChannel chan *proto.DataStoreConnectionTestRequest
 
 	lastTriggerResponse *proto.TriggerResponse
 	lastPollingResponse *proto.PollingResponse
 }
 
 func NewGrpcServer() *GrpcServerMock {
 	server := &GrpcServerMock{
-		triggerChannel:     make(chan *proto.TriggerRequest),
-		pollingChannel:     make(chan *proto.PollingRequest),
-		terminationChannel: make(chan *proto.ShutdownRequest),
+		triggerChannel:       make(chan *proto.TriggerRequest),
+		pollingChannel:       make(chan *proto.PollingRequest),
+		terminationChannel:   make(chan *proto.ShutdownRequest),
+		dataStoreTestChannel: make(chan *proto.DataStoreConnectionTestRequest),
 	}
 	var wg sync.WaitGroup
 	wg.Add(1)
@@ -108,6 +110,20 @@ func (s *GrpcServerMock) RegisterPollerAgent(id *proto.AgentIdentification, stre
 	}
 }
 
+func (s *GrpcServerMock) RegisterDataStoreConnectionTestAgent(id *proto.AgentIdentification, stream proto.Orchestrator_RegisterDataStoreConnectionTestAgentServer) error {
+	if id.Token != "token" {
+		return fmt.Errorf("could not validate token")
+	}
+
+	for {
+		dsTestRequest := <-s.dataStoreTestChannel
+		err := stream.Send(dsTestRequest)
+		if err != nil {
+			log.Println("could not send polling request to agent: %w", err)
+		}
+	}
+}
+
 func (s *GrpcServerMock) SendPolledSpans(ctx context.Context, result *proto.PollingResponse) (*proto.Empty, error) {
 	if result.AgentIdentification == nil || result.AgentIdentification.Token != "token" {
 		return nil, fmt.Errorf("could not validate token")

agent/client/workflow_listen_for_ds_connection_tests.go

@@ -0,0 +1,37 @@
+package client
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"log"
+
+	"github.com/kubeshop/tracetest/agent/proto"
+)
+
+func (c *Client) startDataStoreConnectionTestListener(ctx context.Context) error {
+	client := proto.NewOrchestratorClient(c.conn)
+
+	stream, err := client.RegisterDataStoreConnectionTestAgent(ctx, c.sessionConfig.AgentIdentification)
+	if err != nil {
+		return fmt.Errorf("could not open agent stream: %w", err)
+	}
+
+	go func() {
+		for {
+			req := proto.DataStoreConnectionTestRequest{}
+			err := stream.RecvMsg(&req)
+			if err == io.EOF {
+				return
+			}
+
+			if err != nil {
+				log.Fatal("could not get message from trigger stream: %w", err)
+			}
+
+			// TODO: Get ctx from request
+			c.dataStoreConnectionListener(context.Background(), &req)
+		}
+	}()
+	return nil
+}

agent/client/workflow_send_ds_connection_result.go

@@ -0,0 +1,21 @@
+package client
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/kubeshop/tracetest/agent/proto"
+)
+
+func (c *Client) SendDataStoreConnectionResult(ctx context.Context, response *proto.DataStoreConnectionTestResponse) error {
+	client := proto.NewOrchestratorClient(c.conn)
+
+	response.AgentIdentification = c.sessionConfig.AgentIdentification
+
+	_, err := client.SendDataStoreConnectionTestResult(ctx, response)
+	if err != nil {
+		return fmt.Errorf("could not send data store connection result request: %w", err)
+	}
+
+	return nil
+}

agent/initialization/start.go

@@ -21,7 +21,9 @@ func NewClient(ctx context.Context, config config.Config) (*client.Client, error
 
 	triggerWorker := workers.NewTriggerWorker(client)
 	pollingWorker := workers.NewPollerWorker(client)
+	dataStoreTestConnectionWorker := workers.NewTestConnectionWorker(client)
 
+	client.OnDataStoreTestConnectionRequest(dataStoreTestConnectionWorker.Test)
 	client.OnTriggerRequest(triggerWorker.Trigger)
 	client.OnPollingRequest(pollingWorker.Poll)
 	client.OnConnectionClosed(func(ctx context.Context, sr *proto.ShutdownRequest) error {