go reports-service: Add tests to openai client

Magpie-Monitor · Nov 28, 2024 · f7b3caa · f7b3caa
1 parent 336bd7a
commit f7b3caa
Show file tree

Hide file tree

Showing 3 changed files with 144 additions and 27 deletions.
diff --git a/docker-compose.test.yml b/docker-compose.test.yml
@@ -1,29 +1,29 @@
 services:
-  # logs-ingestion-service-test:
-  #   user: "0"
-  #   container_name: magpie-monitor-test-logs-ingestion-service
-  #   build:
-  #     context: ./go
-  #     dockerfile: ./docker/logs_ingestion/Dockerfile
-  #     target: tests
-  #   environment:
-  #     - APP_ENV=test
-  #     - LOGSDB_PORT=${LOGSDB_PORT}
-  #     - LOGSDB_HOST=${LOGSDB_HOST}
-  #     - LOGSDB_USER=${LOGSDB_USER}
-  #     - LOGSDB_PASSWORD=${LOGSDB_PASSWORD}
-  #     - LOGSDB_CERT_PATH=${LOGSDB_CERT_PATH}
-  #     - LOGS_INGESTION_QUEUE_HOST=${LOGS_INGESTION_QUEUE_HOST}
-  #     - LOGS_INGESTION_QUEUE_PORT=${LOGS_INGESTION_QUEUE_PORT}
-  #     - LOGS_INGESTION_QUEUE_USERNAME=${KAFKA_CLIENT_USERNAME}
-  #     - LOGS_INGESTION_QUEUE_PASSWORD=${KAFKA_CLIENT_PASSWORD}
-  #     - LOGS_INGESTION_NODE_LOGS_TOPIC=${LOGS_INGESTION_NODE_LOGS_TOPIC}
-  #     - LOGS_INGESTION_APPLICATION_LOGS_TOPIC=${LOGS_INGESTION_APPLICATION_LOGS_TOPIC}
-  #     - KAFKA_BROKER_GROUP_ID=${KAFKA_BROKER_GROUP_ID}
-  #     - INTEGRATION_TEST_WAIT_MODIFIER=${INTEGRATION_TEST_WAIT_MODIFIER}
-  #   volumes:
-  #     - es-certs:/usr/local/share
-  #     - cache:/go/pkg/mod/
+  logs-ingestion-service-test:
+    user: "0"
+    container_name: magpie-monitor-test-logs-ingestion-service
+    build:
+      context: ./go
+      dockerfile: ./docker/logs_ingestion/Dockerfile
+      target: tests
+    environment:
+      - APP_ENV=test
+      - LOGSDB_PORT=${LOGSDB_PORT}
+      - LOGSDB_HOST=${LOGSDB_HOST}
+      - LOGSDB_USER=${LOGSDB_USER}
+      - LOGSDB_PASSWORD=${LOGSDB_PASSWORD}
+      - LOGSDB_CERT_PATH=${LOGSDB_CERT_PATH}
+      - LOGS_INGESTION_QUEUE_HOST=${LOGS_INGESTION_QUEUE_HOST}
+      - LOGS_INGESTION_QUEUE_PORT=${LOGS_INGESTION_QUEUE_PORT}
+      - LOGS_INGESTION_QUEUE_USERNAME=${KAFKA_CLIENT_USERNAME}
+      - LOGS_INGESTION_QUEUE_PASSWORD=${KAFKA_CLIENT_PASSWORD}
+      - LOGS_INGESTION_NODE_LOGS_TOPIC=${LOGS_INGESTION_NODE_LOGS_TOPIC}
+      - LOGS_INGESTION_APPLICATION_LOGS_TOPIC=${LOGS_INGESTION_APPLICATION_LOGS_TOPIC}
+      - KAFKA_BROKER_GROUP_ID=${KAFKA_BROKER_GROUP_ID}
+      - INTEGRATION_TEST_WAIT_MODIFIER=${INTEGRATION_TEST_WAIT_MODIFIER}
+    volumes:
+      - es-certs:/usr/local/share
+      - cache:/go/pkg/mod/
   reports-service-test:
     user: "0"
     container_name: magpie-monitor-test-reports-service

diff --git a/go/services/reports/pkg/openai/client.go b/go/services/reports/pkg/openai/client.go
@@ -429,8 +429,6 @@ func (c *Client) SplitCompletionReqestsByBatchSize(completionRequests map[string
 			return nil, err
 		}
 
-		c.logger.Debug("Packet size", zap.Any("len", lastPacketSize))
-		c.logger.Debug("Current size", zap.Any("len", len(encodedPacket)))
 		if lastPacketSize+len(encodedPacket) > c.BatchSizeBytes {
 			requestPackets = append(requestPackets, lastPacket)
 			lastPacket = map[string]*CompletionRequest{customId: request}

diff --git a/go/services/reports/pkg/openai/client_test.go b/go/services/reports/pkg/openai/client_test.go
@@ -0,0 +1,119 @@
+package openai_test
+
+import (
+	"github.com/Magpie-Monitor/magpie-monitor/pkg/tests"
+	"github.com/Magpie-Monitor/magpie-monitor/services/reports/pkg/config"
+	"github.com/Magpie-Monitor/magpie-monitor/services/reports/pkg/openai"
+	"github.com/stretchr/testify/assert"
+	"go.uber.org/fx"
+	"go.uber.org/zap"
+	"testing"
+)
+
+func TestClientSplitCompletionReqestsByBatchSize(t *testing.T) {
+
+	type TestDependencies struct {
+		fx.In
+		Logger *zap.Logger
+		Client *openai.Client
+	}
+
+	testCases := []struct {
+		completionRequests                 map[string]*openai.CompletionRequest
+		batchSizeBytes                     int
+		expectedSplittedCompletionRequests []map[string]*openai.CompletionRequest
+	}{
+		{
+			batchSizeBytes: 180,
+			completionRequests: map[string]*openai.CompletionRequest{
+				"app-1": {
+					// Total length of 171 bytes
+					Messages: []*openai.Message{
+						{
+							Content: "This is the first message content",
+						},
+						{
+							Content: "This is a second message",
+						},
+					},
+				},
+				"app-2": {
+					// Total length of 171 bytes
+					Messages: []*openai.Message{
+						{
+							Content: "This is the first message content",
+						},
+						{
+							Content: "This is a second message",
+						},
+					},
+				},
+				"app-3": {
+					// Total length of 171 bytes
+					Messages: []*openai.Message{
+						{
+							Content: "This is the first message content",
+						},
+						{
+							Content: "This is a second message",
+						},
+					},
+				},
+			},
+			expectedSplittedCompletionRequests: []map[string]*openai.CompletionRequest{
+				{
+					"app-1": {
+						// Total length of 171 bytes
+						Messages: []*openai.Message{
+							{
+								Content: "This is the first message content",
+							},
+							{
+								Content: "This is a second message",
+							},
+						},
+					},
+				},
+				{
+					"app-2": {
+						// Total length of 171 bytes
+						Messages: []*openai.Message{
+							{
+								Content: "This is the first message content",
+							},
+							{
+								Content: "This is a second message",
+							},
+						},
+					},
+				},
+				{
+
+					"app-3": {
+						// Total length of 171 bytes
+						Messages: []*openai.Message{
+							{
+								Content: "This is the first message content",
+							},
+							{
+								Content: "This is a second message",
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+
+	test := func(dependencies TestDependencies) {
+		for _, tc := range testCases {
+			dependencies.Client.BatchSizeBytes = tc.batchSizeBytes
+			splitted, err := dependencies.Client.SplitCompletionReqestsByBatchSize(tc.completionRequests)
+			assert.NoError(t, err, "Failed to split completion requests")
+			assert.Equal(t, tc.expectedSplittedCompletionRequests, splitted)
+
+		}
+	}
+
+	tests.RunTest(test, t, config.AppModule)
+}