From e651f3c3bd779f8c4ee540350abdf022d8eabaab Mon Sep 17 00:00:00 2001
From: Wesley Pettit <wppttt@amazon.com>
Date: Mon, 18 Apr 2022 14:44:07 -0700
Subject: [PATCH] test: add tests for partial metadata flags

Signed-off-by: Wesley Pettit <wppttt@amazon.com>
---
 logger/common_test.go | 59 ++++++++++++++++++++++++++++++++-----------
 1 file changed, 44 insertions(+), 15 deletions(-)

diff --git a/logger/common_test.go b/logger/common_test.go
index 6940c99..6eb8360 100644
--- a/logger/common_test.go
+++ b/logger/common_test.go
@@ -19,6 +19,7 @@ import (
 	"bytes"
 	"bufio"
 	"context"
+	"encoding/json"
 	"io/ioutil"
 	"os"
 	"testing"
@@ -50,13 +51,16 @@ var (
 // TestSendLogs test case as we do not need the functionality that the actual Log function
 // is doing inside the test. Mock Log function is not enough here as there does not exist a
 // better way to verify what happened in the TestSendLogs test, which has a goroutine.
-type dummyClient struct{}
+type dummyClient struct{
+	t *testing.T
+}
 
 // Log implements customized workflow used for testing purpose.
 // This is only trigger in TestSendLogs test case. It writes current log message to the end of
 // tmp test file, which makes sure the function itself accepts and "logging" the message
 // correctly.
 func (d *dummyClient) Log(msg *dockerlogger.Message) error {
+	var b []byte
 	_, err := os.Stat(logDestinationFileName)
 	if err != nil {
 		return err
@@ -67,13 +71,22 @@ func (d *dummyClient) Log(msg *dockerlogger.Message) error {
 			"unable to open file %s to record log message", logDestinationFileName)
 	}
 	defer f.Close()
-	f.Write(msg.Line)
+	b, err = json.Marshal(msg)
+	require.NoError(d.t, err)
+	f.Write(b)
 	f.Write([]byte{'\n'})
 
 	return nil
 }
 
-func checkLogFile(t *testing.T, fileName string, expectedNumLines int) {
+func checkLogFile(t *testing.T, fileName string, expectedNumLines int,
+	              expectedPartialOrdinalSequence[]int) {
+	var (
+		msg dockerlogger.Message
+		line string
+		lastPartialID string
+		lastPartialOrdinal int
+	)
     file, err := os.Open(fileName)
     require.NoError(t, err)
     defer file.Close()
@@ -81,6 +94,21 @@ func checkLogFile(t *testing.T, fileName string, expectedNumLines int) {
     scanner := bufio.NewScanner(file)
     lines := 0
     for scanner.Scan() {
+		line = scanner.Text()
+		err = json.Unmarshal([]byte(line), &msg)
+		require.NoError(t, err)
+		if len(expectedPartialOrdinalSequence) > 0 && lines < len(expectedPartialOrdinalSequence) {
+			// check partial fields
+			require.Equal(t, expectedPartialOrdinalSequence[lines], msg.PLogMetaData.Ordinal)
+			if msg.PLogMetaData.Ordinal < lastPartialOrdinal {
+				// new split message so new partial ID
+				require.NotEqual(t, lastPartialID, msg.PLogMetaData.ID)
+			} else if msg.PLogMetaData.Ordinal > 1 {
+				// this partial ID should be same as last ID
+				require.Equal(t, lastPartialID, msg.PLogMetaData.ID)
+			}
+			lastPartialID = msg.PLogMetaData.ID
+		} 
         lines++
     }
 	require.Equal(t, expectedNumLines, lines)
@@ -95,11 +123,12 @@ func checkLogFile(t *testing.T, fileName string, expectedNumLines int) {
 func TestSendLogs(t *testing.T) {
 
 	for _, tc := range []struct {
-		testName           string
-		bufferSizeInBytes  int
-		maxReadBytes       int
-		logMessages        []string
-		expectedNumOfLines int
+		testName                       string
+		bufferSizeInBytes              int
+		maxReadBytes                   int
+		logMessages                    []string
+		expectedNumOfLines             int
+		expectedPartialOrdinalSequence []int
 	}{
 		{
 			testName:          "general case",
@@ -110,6 +139,7 @@ func TestSendLogs(t *testing.T) {
 				"Second line to write",
 			},
 			expectedNumOfLines: 2, // 2 messages stay as 2 messages
+			expectedPartialOrdinalSequence: []int{}, // neither will be partial
 		},
 		{
 			testName:          "long log message",
@@ -119,6 +149,7 @@ func TestSendLogs(t *testing.T) {
 				"First line to write", // Larger than buffer size.
 			},
 			expectedNumOfLines: 3, // One line 19 chars with 8 char buffer becomes 3 split messages
+			expectedPartialOrdinalSequence: []int{1, 2, 3},
 		},
 		{
 			testName:          "two long log messages",
@@ -129,12 +160,13 @@ func TestSendLogs(t *testing.T) {
 				"Second line to write", // 20 chars => 3 messages
 			},
 			expectedNumOfLines: 6, // 3 + 3 = 6 total
+			expectedPartialOrdinalSequence: []int{1, 2, 3, 1, 2, 3},
 		},
 	} {
 		t.Run(tc.testName, func(t *testing.T) {
 			l := &Logger{
 				Info:              &dockerlogger.Info{},
-				Stream:            &dummyClient{},
+				Stream:            &dummyClient{t},
 				bufferSizeInBytes: tc.bufferSizeInBytes,
 				maxReadBytes:      tc.maxReadBytes,
 			}
@@ -144,21 +176,18 @@ func TestSendLogs(t *testing.T) {
 			require.NoError(t, err)
 			defer os.Remove(tmpIOSource.Name())
 			var (
-				expectedSize int64
 				testPipe     bytes.Buffer
 			)
 			for _, logMessage := range tc.logMessages {
-				expectedSize += int64(len([]rune(logMessage)))
 				_, err := testPipe.WriteString(logMessage + "\n")
 				require.NoError(t, err)
 			}
-			expectedSize += int64(tc.expectedNumOfLines) // for newlines
 
 			// Create a tmp file that used to inside customized dummy Log function where the
 			// logger sends log messages to.
 			tmpDest, err := ioutil.TempFile(os.TempDir(), "")
 			require.NoError(t, err)
-			defer os.Remove(tmpDest.Name())
+			//defer os.Remove(tmpDest.Name())
 			logDestinationFileName = tmpDest.Name()
 
 			var errGroup errgroup.Group
@@ -172,9 +201,9 @@ func TestSendLogs(t *testing.T) {
 			// goroutine.
 			logDestinationInfo, err := os.Stat(logDestinationFileName)
 			require.NoError(t, err)
-			require.Equal(t, expectedSize, logDestinationInfo.Size())
+			require.NotZero(t, logDestinationInfo.Size())
 
-			checkLogFile(t, logDestinationFileName, tc.expectedNumOfLines)
+			checkLogFile(t, logDestinationFileName, tc.expectedNumOfLines, tc.expectedPartialOrdinalSequence)
 		})
 	}
 }