Skip to content

Commit 4c5a316

Browse files
spdeepakspdeepak
committed
add tests for sankey exporter
1 parent 559dbca commit 4c5a316

4 files changed

Lines changed: 209 additions & 48 deletions

File tree

‎exporters/common_test.go‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
package exporters
2+
3+
import (
4+
"context"
5+
"net/http"
6+
"time"
7+
8+
"github.com/spdeepak/flowtracker"
9+
)
10+
11+
func Handler(w http.ResponseWriter, r *http.Request) {
12+
ctx := r.Context()
13+
14+
// 1. Simulate some logic
15+
processPayment(ctx)
16+
17+
// 2. Simulate a Database call
18+
fetchUserProfile(ctx, "user_123")
19+
20+
// 3. Simulate external Microservice call
21+
callShippingService(ctx)
22+
23+
w.Write([]byte("Order Processed"))
24+
}
25+
26+
func processPayment(ctx context.Context) {
27+
// Start a sub-span
28+
ctx, finish := flowtracker.StartSpan(ctx, "Process Payment")
29+
defer finish()
30+
31+
time.Sleep(50 * time.Millisecond) // Simulate work
32+
flowtracker.AddTag(ctx, "currency", "USD")
33+
}
34+
35+
func fetchUserProfile(ctx context.Context, userID string) {
36+
ctx, finish := flowtracker.StartSpan(ctx, "DB: Select User")
37+
defer finish()
38+
39+
flowtracker.AddTag(ctx, "db.query", "SELECT * FROM users...")
40+
time.Sleep(120 * time.Millisecond) // Simulate DB latency
41+
}
42+
43+
func callShippingService(ctx context.Context) {
44+
ctx, finish := flowtracker.StartSpan(ctx, "HTTP: Shipping Service")
45+
defer finish()
46+
47+
// Nested span example: Logic inside the external call preparation
48+
func(innerCtx context.Context) {
49+
_, end := flowtracker.StartSpan(innerCtx, "Calculate Weight")
50+
defer end()
51+
time.Sleep(10 * time.Millisecond)
52+
}(ctx)
53+
54+
time.Sleep(200 * time.Millisecond) // Simulate network call
55+
}

‎exporters/mermaid_test.go‎

Lines changed: 0 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@ package exporters
22

33
import (
44
"bytes"
5-
"context"
65
"io"
76
"net/http"
87
"net/http/httptest"
@@ -158,49 +157,3 @@ func TestMermaidExporter_OK_IncludeTags(t *testing.T) {
158157
t.Fatalf("expected log not found: %s", logs)
159158
}
160159
}
161-
162-
func Handler(w http.ResponseWriter, r *http.Request) {
163-
ctx := r.Context()
164-
165-
// 1. Simulate some logic
166-
processPayment(ctx)
167-
168-
// 2. Simulate a Database call
169-
fetchUserProfile(ctx, "user_123")
170-
171-
// 3. Simulate external Microservice call
172-
callShippingService(ctx)
173-
174-
w.Write([]byte("Order Processed"))
175-
}
176-
177-
func processPayment(ctx context.Context) {
178-
// Start a sub-span
179-
ctx, finish := flowtracker.StartSpan(ctx, "Process Payment")
180-
defer finish()
181-
182-
time.Sleep(50 * time.Millisecond) // Simulate work
183-
flowtracker.AddTag(ctx, "currency", "USD")
184-
}
185-
186-
func fetchUserProfile(ctx context.Context, userID string) {
187-
ctx, finish := flowtracker.StartSpan(ctx, "DB: Select User")
188-
defer finish()
189-
190-
flowtracker.AddTag(ctx, "db.query", "SELECT * FROM users...")
191-
time.Sleep(120 * time.Millisecond) // Simulate DB latency
192-
}
193-
194-
func callShippingService(ctx context.Context) {
195-
ctx, finish := flowtracker.StartSpan(ctx, "HTTP: Shipping Service")
196-
defer finish()
197-
198-
// Nested span example: Logic inside the external call preparation
199-
func(innerCtx context.Context) {
200-
_, end := flowtracker.StartSpan(innerCtx, "Calculate Weight")
201-
defer end()
202-
time.Sleep(10 * time.Millisecond)
203-
}(ctx)
204-
205-
time.Sleep(200 * time.Millisecond) // Simulate network call
206-
}

‎exporters/sankey.go‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package exporters
22

33
import (
44
"fmt"
5+
"sort"
56
"strings"
67

78
"github.com/spdeepak/flowtracker"
@@ -16,6 +17,9 @@ type SankeyExporter struct {
1617
// List of tag keys you want to display in the diagram.
1718
// Example: []string{"db.table", "http.status_code"}
1819
IncludeTags []string
20+
21+
// IncludeAllTags overrides IncludeTags. If true, ALL tags present
22+
IncludeAllTags bool
1923
}
2024

2125
func (s *SankeyExporter) Export(tr *flowtracker.Trace) {
@@ -30,7 +34,15 @@ func (s *SankeyExporter) Export(tr *flowtracker.Trace) {
3034

3135
// Check if user wants to see tags for this span
3236
var tagSuffixes []string
33-
if len(s.IncludeTags) > 0 && span.Tags != nil {
37+
// Logic: Decide which keys to show
38+
if s.IncludeAllTags {
39+
// Get ALL keys from the map
40+
for key, val := range span.Tags {
41+
tagSuffixes = append(tagSuffixes, fmt.Sprintf("%s:%s", key, val))
42+
}
43+
// Sort keys to ensure deterministic diagram output
44+
sort.Strings(tagSuffixes)
45+
} else if len(s.IncludeTags) > 0 && span.Tags != nil {
3446
for _, key := range s.IncludeTags {
3547
if val, ok := span.Tags[key]; ok {
3648
// Format: (key:value)

‎exporters/sankey_test.go‎

Lines changed: 141 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,141 @@
1+
package exporters
2+
3+
import (
4+
"bytes"
5+
"io"
6+
"net/http"
7+
"net/http/httptest"
8+
"os"
9+
"strings"
10+
"testing"
11+
"time"
12+
13+
"github.com/spdeepak/flowtracker"
14+
)
15+
16+
func NewSankeyServer(tags []string, includeAllTags bool) http.Handler {
17+
mux := http.NewServeMux()
18+
mux.HandleFunc("/", Handler)
19+
mw := flowtracker.NewMiddleware(flowtracker.WithExporter(&SankeyExporter{IncludeTags: tags, IncludeAllTags: includeAllTags}))
20+
return mw(mux)
21+
}
22+
23+
func TestSankeyExporter_OK(t *testing.T) {
24+
// capture stdout
25+
old := os.Stdout
26+
r, w, _ := os.Pipe()
27+
os.Stdout = w
28+
29+
server := NewSankeyServer([]string{}, true)
30+
31+
req := httptest.NewRequest(http.MethodGet, "/", nil)
32+
rr := httptest.NewRecorder()
33+
34+
server.ServeHTTP(rr, req)
35+
36+
time.Sleep(1 * time.Second)
37+
38+
// restore stdout
39+
w.Close()
40+
os.Stdout = old
41+
42+
// read logs
43+
var buf bytes.Buffer
44+
_, err := io.Copy(&buf, r)
45+
if err != nil {
46+
t.Error("Error should be nil")
47+
}
48+
49+
logs := buf.String()
50+
output := strings.Split(logs, "\n")
51+
if !strings.Contains(output[1], "----- START SANKEY DATA (trace id: trace-") {
52+
t.Fatalf("expected log not found: %s", output[1])
53+
}
54+
if !strings.Contains(output[2], "GET / [") {
55+
t.Fatalf("expected log not found: %s", output[2])
56+
}
57+
if !strings.Contains(output[2], "] Process Payment (currency:USD)") {
58+
t.Fatalf("expected log not found: %s", output[2])
59+
}
60+
if !strings.Contains(output[3], "GET / [") {
61+
t.Fatalf("expected log not found: %s", output[3])
62+
}
63+
if !strings.Contains(output[3], "] DB: Select User (db.query:SELECT * FROM users...)") {
64+
t.Fatalf("expected log not found: %s", output[3])
65+
}
66+
if !strings.Contains(output[4], "GET / [") {
67+
t.Fatalf("expected log not found: %s", output[4])
68+
}
69+
if !strings.Contains(output[4], "] HTTP: Shipping Service") {
70+
t.Fatalf("expected log not found: %s", output[4])
71+
}
72+
if !strings.Contains(output[5], "HTTP: Shipping Service [") {
73+
t.Fatalf("expected log not found: %s", output[5])
74+
}
75+
if !strings.Contains(output[5], "] Calculate Weight") {
76+
t.Fatalf("expected log not found: %s", output[5])
77+
}
78+
if !strings.Contains(output[6], "----- END SANKEY DATA (trace id: trace-") {
79+
t.Fatalf("expected log not found: %s", output[6])
80+
}
81+
}
82+
83+
func TestSankeyExporter_OK_IncludeTags(t *testing.T) {
84+
// capture stdout
85+
old := os.Stdout
86+
r, w, _ := os.Pipe()
87+
os.Stdout = w
88+
89+
server := NewSankeyServer([]string{"db.query"}, false)
90+
91+
req := httptest.NewRequest(http.MethodGet, "/", nil)
92+
rr := httptest.NewRecorder()
93+
94+
server.ServeHTTP(rr, req)
95+
96+
time.Sleep(1 * time.Second)
97+
98+
// restore stdout
99+
w.Close()
100+
os.Stdout = old
101+
102+
// read logs
103+
var buf bytes.Buffer
104+
_, err := io.Copy(&buf, r)
105+
if err != nil {
106+
t.Error("Error should be nil")
107+
}
108+
109+
logs := buf.String()
110+
output := strings.Split(logs, "\n")
111+
if !strings.Contains(output[1], "----- START SANKEY DATA (trace id") {
112+
t.Fatalf("expected log not found: %s", output[1])
113+
}
114+
if !strings.Contains(output[2], "GET / [") {
115+
t.Fatalf("expected log not found: %s", output[2])
116+
}
117+
if !strings.Contains(output[2], "] Process Payment") {
118+
t.Fatalf("expected log not found: %s", output[2])
119+
}
120+
if !strings.Contains(output[3], "GET / [") {
121+
t.Fatalf("expected log not found: %s", output[3])
122+
}
123+
if !strings.Contains(output[3], "] DB: Select User (db.query:SELECT * FROM users...)") {
124+
t.Fatalf("expected log not found: %s", output[3])
125+
}
126+
if !strings.Contains(output[4], "GET / [") {
127+
t.Fatalf("expected log not found: %s", output[4])
128+
}
129+
if !strings.Contains(output[4], "] HTTP: Shipping Service") {
130+
t.Fatalf("expected log not found: %s", output[4])
131+
}
132+
if !strings.Contains(output[5], "HTTP: Shipping Service [") {
133+
t.Fatalf("expected log not found: %s", output[5])
134+
}
135+
if !strings.Contains(output[5], "] Calculate Weight") {
136+
t.Fatalf("expected log not found: %s", output[5])
137+
}
138+
if !strings.Contains(output[6], "----- END SANKEY DATA (trace id: trace-") {
139+
t.Fatalf("expected log not found: %s", output[6])
140+
}
141+
}

0 commit comments

Comments
 (0)