From b957a3959abd05c8e2764f1b63fb88cea13bcfb7 Mon Sep 17 00:00:00 2001
From: Matheus Nogueira <matheus.nogueira2008@gmail.com>
Date: Fri, 15 Sep 2023 17:06:56 -0300
Subject: [PATCH] feat(server): allow trace to be incremented in UpdateTestRun
 endpoint (#3128)

* feat(server): allow trace to be incremented in UpdateTestRun endpoint

* fix: make sure trace.Spans() consider tree object when flat is nil

* feat(server): allow trace to be incremented in UpdateTestRun endpoint

* fix: make sure trace.Spans() consider tree object when flat is nil

* fix: panic when adding attributes to span

* remove condition
---
 server/http/controller.go       |  3 +++
 server/traces/span_entitiess.go |  4 ++++
 server/traces/trace_entities.go | 18 ++++++++++++++++++
 server/traces/traces_test.go    | 19 +++++++++++++++++++
 4 files changed, 44 insertions(+)

diff --git a/server/http/controller.go b/server/http/controller.go
index fab9cd0359..c3a5796567 100644
--- a/server/http/controller.go
+++ b/server/http/controller.go
@@ -22,6 +22,7 @@ import (
 	"github.com/kubeshop/tracetest/server/test"
 	"github.com/kubeshop/tracetest/server/testsuite"
 	"github.com/kubeshop/tracetest/server/tracedb"
+	"github.com/kubeshop/tracetest/server/traces"
 	"github.com/kubeshop/tracetest/server/variableset"
 	"github.com/labstack/gommon/log"
 	"go.opentelemetry.io/otel/trace"
@@ -718,6 +719,8 @@ func (c *controller) UpdateTestRun(ctx context.Context, testID string, runID int
 
 	// Prevents bad data in other fields to override correct data
 	existingRun.TriggerResult = run.TriggerResult
+	newTrace := traces.MergeTraces(existingRun.Trace, run.Trace)
+	existingRun.Trace = &newTrace
 
 	err = c.testRunRepository.UpdateRun(ctx, existingRun)
 	if err != nil {
diff --git a/server/traces/span_entitiess.go b/server/traces/span_entitiess.go
index d907b2b713..b630ce387d 100644
--- a/server/traces/span_entitiess.go
+++ b/server/traces/span_entitiess.go
@@ -226,6 +226,10 @@ func decodeChildren(parent *Span, children []encodedSpan, cache spanCache) ([]*S
 }
 
 func (span Span) setMetadataAttributes() Span {
+	if span.Attributes == nil {
+		span.Attributes = Attributes{}
+	}
+
 	span.Attributes[TracetestMetadataFieldName] = span.Name
 	span.Attributes[TracetestMetadataFieldType] = spanType(span.Attributes)
 	span.Attributes[TracetestMetadataFieldDuration] = spanDuration(span)
diff --git a/server/traces/trace_entities.go b/server/traces/trace_entities.go
index b668ca106d..8f8650c294 100644
--- a/server/traces/trace_entities.go
+++ b/server/traces/trace_entities.go
@@ -18,6 +18,24 @@ type Trace struct {
 	Flat     map[trace.SpanID]*Span `json:"-"`
 }
 
+func MergeTraces(traces ...*Trace) Trace {
+	if len(traces) == 0 {
+		return NewTrace(id.NewRandGenerator().TraceID().String(), []Span{})
+	}
+
+	traceID := traces[0].ID
+	spans := make([]Span, 0)
+	for _, trace := range traces {
+		if trace == nil {
+			continue
+		}
+
+		spans = append(spans, trace.Spans()...)
+	}
+
+	return NewTrace(traceID.String(), spans)
+}
+
 func NewTrace(traceID string, spans []Span) Trace {
 	spanMap := make(map[string]*Span, 0)
 	for _, span := range spans {
diff --git a/server/traces/traces_test.go b/server/traces/traces_test.go
index 124892930d..5d55a0dddb 100644
--- a/server/traces/traces_test.go
+++ b/server/traces/traces_test.go
@@ -204,6 +204,25 @@ func TestEventsAreInjectedIntoAttributes(t *testing.T) {
 	assert.Equal(t, "value", events[1].Attributes["attribute2"])
 }
 
+func TestMergingZeroTraces(t *testing.T) {
+	trace := traces.MergeTraces()
+	assert.NotEmpty(t, trace.ID)
+}
+
+func TestMergingFragmentsFromSameTrace(t *testing.T) {
+	traceID := id.NewRandGenerator().TraceID().String()
+	rootSpan := newSpan("Root")
+	childSpan1 := newSpan("child 1", withParent(&rootSpan))
+	childSpan2 := newSpan("child 2", withParent(&rootSpan))
+
+	firstFragment := traces.NewTrace(traceID, []traces.Span{childSpan2})
+	secondFragment := traces.NewTrace(traceID, []traces.Span{rootSpan, childSpan1})
+
+	trace := traces.MergeTraces(&firstFragment, &secondFragment)
+	assert.NotEmpty(t, trace.ID)
+	assert.Len(t, trace.Flat, 3)
+}
+
 type option func(*traces.Span)
 
 func withParent(parent *traces.Span) option {