From 6fcdf2fd9851d8f7edb93365788a307b0ffde6bd Mon Sep 17 00:00:00 2001
From: Ravshan Yuldashev <38344113+rorua@users.noreply.github.com>
Date: Mon, 11 Aug 2025 10:57:45 +0500
Subject: [PATCH] Add tests for baggage helpers

---
 baggage.go      |  8 ++++----
 baggage_test.go | 29 +++++++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 4 deletions(-)
 create mode 100644 baggage_test.go

diff --git a/baggage.go b/baggage.go
index dc6e2e2..f36e185 100644
--- a/baggage.go
+++ b/baggage.go
@@ -13,19 +13,19 @@ func GetBaggage(ctx context.Context) baggage.Baggage {
 
 // AddBaggageItem adds a key-value pair to the baggage.
 func AddBaggageItem(ctx context.Context, key, value string) context.Context {
+	b := baggage.FromContext(ctx)
 	m, _ := baggage.NewMember(key, value)
-	b, _ := baggage.New(m)
+	b, _ = b.SetMember(m)
 	return baggage.ContextWithBaggage(ctx, b)
 }
 
 // AddBaggageItems adds multiple key-value pairs to the baggage.
 func AddBaggageItems(ctx context.Context, items map[string]string) context.Context {
-	var members []baggage.Member
+	b := baggage.FromContext(ctx)
 	for key, value := range items {
 		m, _ := baggage.NewMember(key, value)
-		members = append(members, m)
+		b, _ = b.SetMember(m)
 	}
-	b, _ := baggage.New(members...)
 	return baggage.ContextWithBaggage(ctx, b)
 }
 
diff --git a/baggage_test.go b/baggage_test.go
new file mode 100644
index 0000000..04731d7
--- /dev/null
+++ b/baggage_test.go
@@ -0,0 +1,29 @@
+package otelemetry
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestAddBaggageItem(t *testing.T) {
+	ctx := context.Background()
+	ctx = AddBaggageItem(ctx, "user", "alice")
+	ctx = AddBaggageItem(ctx, "tenant", "acme")
+
+	b := GetBaggage(ctx)
+	assert.Equal(t, "alice", b.Member("user").Value())
+	assert.Equal(t, "acme", b.Member("tenant").Value())
+}
+
+func TestAddBaggageItems(t *testing.T) {
+	ctx := context.Background()
+	ctx = AddBaggageItem(ctx, "user", "alice")
+	ctx = AddBaggageItems(ctx, map[string]string{"tenant": "acme", "role": "admin"})
+
+	b := GetBaggage(ctx)
+	assert.Equal(t, "alice", b.Member("user").Value())
+	assert.Equal(t, "acme", b.Member("tenant").Value())
+	assert.Equal(t, "admin", b.Member("role").Value())
+}