Added test coverage for pkg/stream package

Signed-off-by: Shubham Singh <shubhammahar1306@gmail.com>

6 files changed, 486 insertions, 1 deletions

diff --git a/pkg/stream/edgedattachconnection_test.go b/pkg/stream/edgedattachconnection_test.go
new file mode 100644
index 000000000..f7a347105
--- /dev/null
+++ b/pkg/stream/edgedattachconnection_test.go
@@ -0,0 +1,94 @@
+/*
+Copyright 2024 The KubeEdge Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+   http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package stream
+
+import (
+	"encoding/json"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCreateConnectMessage(t *testing.T) {
+	assert := assert.New(t)
+	edgedAttachConn := &EdgedAttachConnection{
+		MessID: 1,
+	}
+
+	msg, err := edgedAttachConn.CreateConnectMessage()
+	assert.NoError(err)
+
+	expectedData, err := json.Marshal(edgedAttachConn)
+	assert.NoError(err)
+	expectedMessage := NewMessage(edgedAttachConn.MessID, MessageTypeAttachConnect, expectedData)
+
+	assert.Equal(expectedMessage, msg)
+}
+
+func TestGetMessageID(t *testing.T) {
+	assert := assert.New(t)
+
+	edgedAttachConn := &EdgedAttachConnection{
+		MessID: uint64(100),
+	}
+
+	messID := edgedAttachConn.GetMessageID()
+	stdResult := uint64(100)
+
+	assert.Equal(messID, stdResult)
+}
+
+func TestString(t *testing.T) {
+	assert := assert.New(t)
+
+	edgedAttachConn := &EdgedAttachConnection{
+		MessID: uint64(100),
+	}
+
+	stdResult := "EDGE_ATTACH_CONNECTOR Message MessageID 100"
+	result := edgedAttachConn.String()
+
+	assert.Equal(result, stdResult)
+}
+
+func TestCacheTunnelMessage(t *testing.T) {
+	assert := assert.New(t)
+	edgedAttachConn := &EdgedAttachConnection{
+		ReadChan: make(chan *Message, 1),
+	}
+
+	msg := &Message{ConnectID: 100, MessageType: MessageTypeData, Data: []byte("test data")}
+	edgedAttachConn.CacheTunnelMessage(msg)
+
+	assert.Equal(msg, <-edgedAttachConn.ReadChan)
+}
+
+func TestCloseReadChannel(t *testing.T) {
+	assert := assert.New(t)
+	edgedAttachConn := &EdgedAttachConnection{
+		ReadChan: make(chan *Message),
+	}
+
+	go func() {
+		time.Sleep(1 * time.Second)
+		edgedAttachConn.CloseReadChannel()
+	}()
+
+	_, ok := <-edgedAttachConn.ReadChan
+	assert.False(ok)
+}
diff --git a/pkg/stream/edgedexecconnection_test.go b/pkg/stream/edgedexecconnection_test.go
new file mode 100644
index 000000000..0d04af82e
--- /dev/null
+++ b/pkg/stream/edgedexecconnection_test.go
@@ -0,0 +1,90 @@
+/*
+Copyright 2024 The KubeEdge Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+   http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package stream
+
+import (
+	"encoding/json"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestExecConnection_CreateConnectMessage(t *testing.T) {
+	assert := assert.New(t)
+
+	edgedExecConn := &EdgedExecConnection{
+		MessID: 1,
+	}
+	msg, err := edgedExecConn.CreateConnectMessage()
+	assert.NoError(err)
+
+	exceptedData, err := json.Marshal(edgedExecConn)
+	assert.NoError(err)
+	expectedMessage := NewMessage(edgedExecConn.MessID, MessageTypeExecConnect, exceptedData)
+
+	assert.Equal(expectedMessage, msg)
+}
+
+func TestExecConnection_GetMessageID(t *testing.T) {
+	assert := assert.New(t)
+	edgedExecConn := &EdgedExecConnection{
+		MessID: uint64(100),
+	}
+
+	stdResult := uint64(100)
+	assert.Equal(stdResult, edgedExecConn.MessID)
+}
+
+func TestExecConnection_String(t *testing.T) {
+	assert := assert.New(t)
+
+	edgedExecConn := &EdgedExecConnection{
+		MessID: uint64(100),
+	}
+
+	result := edgedExecConn.String()
+	stdResult := "EDGE_EXEC_CONNECTOR Message MessageID 100"
+	assert.Equal(stdResult, result)
+}
+
+func TestExecConnection_CacheTunnelMessage(t *testing.T) {
+	assert := assert.New(t)
+	edgedExecConn := &EdgedExecConnection{
+		ReadChan: make(chan *Message, 1),
+	}
+
+	msg := &Message{ConnectID: 100, MessageType: MessageTypeData, Data: []byte("test data")}
+	edgedExecConn.CacheTunnelMessage(msg)
+
+	assert.Equal(msg, <-edgedExecConn.ReadChan)
+}
+
+func TestExecConnection_CloseReadChannel(t *testing.T) {
+	assert := assert.New(t)
+	edgedExecConn := &EdgedExecConnection{
+		ReadChan: make(chan *Message),
+	}
+
+	go func() {
+		time.Sleep(1 * time.Second)
+		edgedExecConn.CloseReadChannel()
+	}()
+
+	_, ok := <-edgedExecConn.ReadChan
+	assert.False(ok)
+}
diff --git a/pkg/stream/edgedlogconnection_test.go b/pkg/stream/edgedlogconnection_test.go
new file mode 100644
index 000000000..9d9e42292
--- /dev/null
+++ b/pkg/stream/edgedlogconnection_test.go
@@ -0,0 +1,90 @@
+/*
+Copyright 2024 The KubeEdge Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+   http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package stream
+
+import (
+	"encoding/json"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestLogConnection_CreateConnectMessage(t *testing.T) {
+	assert := assert.New(t)
+
+	edgedLogsConn := &EdgedLogsConnection{
+		MessID: 1,
+	}
+	msg, err := edgedLogsConn.CreateConnectMessage()
+	assert.NoError(err)
+
+	exceptedData, err := json.Marshal(edgedLogsConn)
+	assert.NoError(err)
+	expectedMessage := NewMessage(edgedLogsConn.MessID, MessageTypeLogsConnect, exceptedData)
+
+	assert.Equal(expectedMessage, msg)
+}
+
+func TestLogConnection_GetMessageID(t *testing.T) {
+	assert := assert.New(t)
+	edgedLogsConn := &EdgedLogsConnection{
+		MessID: uint64(100),
+	}
+
+	stdResult := uint64(100)
+	assert.Equal(stdResult, edgedLogsConn.MessID)
+}
+
+func TestLogConnection_String(t *testing.T) {
+	assert := assert.New(t)
+
+	edgedLogsConn := &EdgedLogsConnection{
+		MessID: uint64(100),
+	}
+
+	result := edgedLogsConn.String()
+	stdResult := "EDGE_LOGS_CONNECTOR Message MessageID 100"
+	assert.Equal(stdResult, result)
+}
+
+func TestLogConnection_CacheTunnelMessage(t *testing.T) {
+	assert := assert.New(t)
+	edgedLogsConn := &EdgedLogsConnection{
+		ReadChan: make(chan *Message, 1),
+	}
+
+	msg := &Message{ConnectID: 100, MessageType: MessageTypeData, Data: []byte("test data")}
+	edgedLogsConn.CacheTunnelMessage(msg)
+
+	assert.Equal(msg, <-edgedLogsConn.ReadChan)
+}
+
+func TestLogConnection_CloseReadChannel(t *testing.T) {
+	assert := assert.New(t)
+	edgedLogsConn := &EdgedLogsConnection{
+		ReadChan: make(chan *Message),
+	}
+
+	go func() {
+		time.Sleep(1 * time.Second)
+		edgedLogsConn.CloseReadChannel()
+	}()
+
+	_, ok := <-edgedLogsConn.ReadChan
+	assert.False(ok)
+}
diff --git a/pkg/stream/edgedmetricsconnection.go b/pkg/stream/edgedmetricsconnection.go
index c34f1818e..55e671fe2 100644
--- a/pkg/stream/edgedmetricsconnection.go
+++ b/pkg/stream/edgedmetricsconnection.go
@@ -65,7 +65,7 @@ func (ms *EdgedMetricsConnection) CreateConnectMessage() (*Message, error) {
 }
 
 func (ms *EdgedMetricsConnection) String() string {
-	return fmt.Sprintf("EDGE_Metrics_CONNECTOR Message MessageID %v", ms.MessID)
+	return fmt.Sprintf("EDGE_METRICS_CONNECTOR Message MessageID %v", ms.MessID)
 }
 
 func (ms *EdgedMetricsConnection) receiveFromCloudStream(stop chan struct{}) {
diff --git a/pkg/stream/edgedmetricsconnection_test.go b/pkg/stream/edgedmetricsconnection_test.go
new file mode 100644
index 000000000..908ac3846
--- /dev/null
+++ b/pkg/stream/edgedmetricsconnection_test.go
@@ -0,0 +1,94 @@
+/*
+Copyright 2024 The KubeEdge Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+   http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package stream
+
+import (
+	"encoding/json"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestMetricsConnection_CreateConnectMessage(t *testing.T) {
+	assert := assert.New(t)
+	edgedMetricsConn := &EdgedMetricsConnection{
+		MessID: 1,
+	}
+
+	msg, err := edgedMetricsConn.CreateConnectMessage()
+	assert.NoError(err)
+
+	expectedData, err := json.Marshal(edgedMetricsConn)
+	assert.NoError(err)
+	expectedMessage := NewMessage(edgedMetricsConn.MessID, MessageTypeMetricConnect, expectedData)
+
+	assert.Equal(expectedMessage, msg)
+}
+
+func TestMetricsConnection_GetMessageID(t *testing.T) {
+	assert := assert.New(t)
+
+	edgedMetricsConn := &EdgedMetricsConnection{
+		MessID: uint64(100),
+	}
+
+	messID := edgedMetricsConn.GetMessageID()
+	stdResult := uint64(100)
+
+	assert.Equal(messID, stdResult)
+}
+
+func TestMetricsConnection_String(t *testing.T) {
+	assert := assert.New(t)
+
+	edgedMetricsConn := &EdgedMetricsConnection{
+		MessID: uint64(100),
+	}
+
+	stdResult := "EDGE_METRICS_CONNECTOR Message MessageID 100"
+	result := edgedMetricsConn.String()
+
+	assert.Equal(result, stdResult)
+}
+
+func TestMetricsConnection_CacheTunnelMessage(t *testing.T) {
+	assert := assert.New(t)
+	edgedMetricsConn := &EdgedMetricsConnection{
+		ReadChan: make(chan *Message, 1),
+	}
+
+	msg := &Message{ConnectID: 100, MessageType: MessageTypeData, Data: []byte("test data")}
+	edgedMetricsConn.CacheTunnelMessage(msg)
+
+	assert.Equal(msg, <-edgedMetricsConn.ReadChan)
+}
+
+func TestMetricsConnection_CloseReadChannel(t *testing.T) {
+	assert := assert.New(t)
+	edgedMetricsConn := &EdgedMetricsConnection{
+		ReadChan: make(chan *Message),
+	}
+
+	go func() {
+		time.Sleep(1 * time.Second)
+		edgedMetricsConn.CloseReadChannel()
+	}()
+
+	_, ok := <-edgedMetricsConn.ReadChan
+	assert.False(ok)
+}
diff --git a/pkg/stream/message_test.go b/pkg/stream/message_test.go
new file mode 100644
index 000000000..2f72d719f
--- /dev/null
+++ b/pkg/stream/message_test.go
@@ -0,0 +1,117 @@
+/*
+Copyright 2024 The KubeEdge Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+   http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package stream
+
+import (
+	"bytes"
+	"encoding/binary"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestMessageType_String(t *testing.T) {
+	assert := assert.New(t)
+	cases := []struct {
+		msg       MessageType
+		stdResult string
+	}{
+		{
+			msg:       MessageTypeLogsConnect,
+			stdResult: "LOGS_CONNECT",
+		},
+		{
+			msg:       MessageTypeExecConnect,
+			stdResult: "EXEC_CONNECT",
+		},
+		{
+			msg:       MessageTypeAttachConnect,
+			stdResult: "ATTACH_CONNECT",
+		},
+		{
+			msg:       MessageTypeMetricConnect,
+			stdResult: "METRIC_CONNECT",
+		},
+		{
+			msg:       MessageTypeData,
+			stdResult: "DATA",
+		},
+		{
+			msg:       MessageTypeRemoveConnect,
+			stdResult: "REMOVE_CONNECT",
+		},
+		{
+			msg:       100,
+			stdResult: "UNKNOWN",
+		},
+	}
+
+	for _, test := range cases {
+		assert.Equal(test.stdResult, test.msg.String())
+	}
+}
+
+func TestMessage_NewMessage(t *testing.T) {
+	assert := assert.New(t)
+	msg := NewMessage(100, MessageTypeLogsConnect, []byte("test data"))
+	assert.Equal(msg.ConnectID, uint64(100))
+	assert.Equal(msg.MessageType, MessageTypeLogsConnect)
+	assert.Equal(msg.Data, []byte("test data"))
+}
+
+func TestMessage_Bytes(t *testing.T) {
+	assert := assert.New(t)
+	message := NewMessage(1, MessageTypeLogsConnect, []byte("test_data"))
+
+	var stdResult []byte
+	buf := make([]byte, 10)
+	n := binary.PutUvarint(buf, message.ConnectID)
+	stdResult = append(stdResult, buf[:n]...)
+	n = binary.PutUvarint(buf, uint64(message.MessageType))
+	stdResult = append(stdResult, buf[:n]...)
+	stdResult = append(stdResult, message.Data...)
+
+	assert.Equal(stdResult, message.Bytes())
+}
+
+func TestMessage_String(t *testing.T) {
+	assert := assert.New(t)
+	msg := &Message{
+		ConnectID:   100,
+		MessageType: MessageTypeLogsConnect,
+		Data:        []byte("test data"),
+	}
+
+	result := msg.String()
+	stdResult := "MESSAGE: connectID 100 messageType LOGS_CONNECT"
+	assert.Equal(result, stdResult)
+}
+
+func TestReadMessageFromTunnel(t *testing.T) {
+	assert := assert.New(t)
+
+	message := NewMessage(100, MessageTypeLogsConnect, []byte("test data"))
+	messageBytes := message.Bytes()
+
+	buffer := bytes.NewBuffer(messageBytes)
+
+	readMessage, err := ReadMessageFromTunnel(buffer)
+	assert.NoError(err)
+	assert.Equal(readMessage.ConnectID, uint64(100))
+	assert.Equal(readMessage.MessageType, MessageTypeLogsConnect)
+	assert.Equal(readMessage.Data, []byte("test data"))
+}