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

@@ -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 {

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)

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 {

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 {