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

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/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 {
@@ -176,6 +194,11 @@ func (t *Trace) Spans() []Span {
 		return []Span{}
 	}
 
+	if len(t.Flat) == 0 {
+		t.Flat = make(map[trace.SpanID]*Span)
+		flattenSpans(t.Flat, t.RootSpan)
+	}
+
 	spans := make([]Span, 0, len(t.Flat))
 	for _, span := range t.Flat {
 		spans = append(spans, *span)

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 {