add unit test for merge

2 files changed, 85 insertions, 2 deletions

diff --git a/internal/flamegraph/iordata.go b/internal/flamegraph/iordata.go
index b6bb197..77b03aa 100644
--- a/internal/flamegraph/iordata.go
+++ b/internal/flamegraph/iordata.go
@@ -53,7 +53,6 @@ func (iod iorData) add(ev *event.Pair) {
 		ev.EnterEv.GetPid(), ev.EnterEv.GetTid(), ev.File.FlagsString(), cnt)
 }
 
-// TODO: Unit test
 func (iod iorData) addPath(path pathType, traceId traceIdType, comm commType,
 	pid pidType, tid tidType, flags flagsType, addCnt counter) {
 
@@ -90,6 +89,38 @@ func (iod iorData) addPath(path pathType, traceId traceIdType, comm commType,
 	}
 }
 
+func (iod iorData) merge(other iorData) iorData {
+	for path, traceIdMap := range other.paths {
+		if _, ok := iod.paths[path]; !ok {
+			iod.paths[path] = make(map[traceIdType]map[commType]map[pidType]map[tidType]map[flagsType]counter)
+		}
+		for traceId, commMap := range traceIdMap {
+			if _, ok := iod.paths[path][traceId]; !ok {
+				iod.paths[path][traceId] = make(map[commType]map[pidType]map[tidType]map[flagsType]counter)
+			}
+			for comm, pidMap := range commMap {
+				if _, ok := iod.paths[path][traceId][comm]; !ok {
+					iod.paths[path][traceId][comm] = make(map[pidType]map[tidType]map[flagsType]counter)
+				}
+				for pid, tidMap := range pidMap {
+					if _, ok := iod.paths[path][traceId][comm][pid]; !ok {
+						iod.paths[path][traceId][comm][pid] = make(map[tidType]map[flagsType]counter)
+					}
+					for tid, flagsMap := range tidMap {
+						if _, ok := iod.paths[path][traceId][comm][pid][tid]; !ok {
+							iod.paths[path][traceId][comm][pid][tid] = make(map[flagsType]counter)
+						}
+						for flags, cnt := range flagsMap {
+							iod.addPath(path, traceId, comm, pid, tid, flags, cnt)
+						}
+					}
+				}
+			}
+		}
+	}
+	return iod
+}
+
 func (iod iorData) commit() error {
 	hostname, err := os.Hostname()
 	if err != nil {
diff --git a/internal/flamegraph/iordata_test.go b/internal/flamegraph/iordata_test.go
index 79fb75c..8a6c5e0 100644
--- a/internal/flamegraph/iordata_test.go
+++ b/internal/flamegraph/iordata_test.go
@@ -1,6 +1,8 @@
 package flamegraph
 
-import "testing"
+import (
+	"testing"
+)
 
 func TestAddPath(t *testing.T) {
 	iod := newIorData()
@@ -26,3 +28,53 @@ func TestAddPath(t *testing.T) {
 		t.Errorf("Expected counter %v, got %v", resultCnt, iod.paths[path][traceId][comm][pid][tid][flags])
 	}
 }
+
+func TestMerge(t *testing.T) {
+	rdwrFlag := flagsType("O_RDWR")
+	roFlag := flagsType("O_RDONLY")
+
+	// Initialize two iorData instances with sample data
+	iod1 := iorData{paths: pathMap{
+		"path1": {1: {"comm1": {100: {1000: {rdwrFlag: counter{
+			count:          10,
+			duration:       1000,
+			durationToPrev: 100,
+		}}}}}}}}
+	iod2 := iorData{paths: pathMap{
+		"path1": {1: {"comm1": {100: {1000: {roFlag: counter{
+			count:          20,
+			duration:       2000,
+			durationToPrev: 200,
+		}}}}}}}}
+	iod3 := iorData{paths: pathMap{
+		"path2": {1: {"comm2": {101: {1000: {roFlag: counter{
+			count:          20,
+			duration:       2000,
+			durationToPrev: 200,
+		}}}}}}}}
+	iod4 := iorData{paths: pathMap{
+		"path2": {1: {"comm2": {101: {1000: {roFlag: counter{
+			count:          40,
+			duration:       4000,
+			durationToPrev: 400,
+		}}}}}}}}
+
+	// Merge iod2 into iod1
+	t.Log(iod1)
+	t.Log(iod2)
+	merged := iod1.merge(iod2).merge(iod3).merge(iod4)
+	t.Log(merged)
+
+	// Check if the merged data contains data from both iod1 and iod2
+	if len(merged.paths) != 2 {
+		t.Errorf("Expected 2 paths, got %d", len(merged.paths))
+	}
+
+	if merged.paths["path1"][1]["comm1"][100][1000][flagsType("O_RDWR")].count != 10 {
+		t.Errorf("Expected counter 10, got %d", merged.paths["path1"][1]["comm1"][100][1000][flagsType("O_RDWR")].count)
+	}
+
+	if merged.paths["path2"][1]["comm2"][101][1000][roFlag].count != 60 {
+		t.Errorf("Expected counter 60, got %d", merged.paths["path2"][1]["comm2"][101][1000][roFlag].count)
+	}
+}