Skip to content

Commit 9bc2fa8

Browse files
sandrolainsandrolain
committed
feat(tools): add support for custom template delimiters
1 parent 61b278d commit 9bc2fa8

9 files changed

Lines changed: 692 additions & 26 deletions

File tree

‎coaptool/send.go‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,9 @@ func sendCommand() *cobra.Command {
2424
sendInterval string
2525
sendProto string
2626
sendMIME string
27+
headers []string
28+
openDelim string
29+
closeDelim string
2730
)
2831

2932
cmd := &cobra.Command{
@@ -36,11 +39,17 @@ func sendCommand() *cobra.Command {
3639
logger := toolutil.Logger()
3740
logger.Info("Sending CoAP POST periodically", "proto", sendProto, "addr", sendAddress, "path", sendPath, "interval", sendInterval)
3841

42+
_, err := toolutil.ParseHeadersWithDelimiters(headers, openDelim, closeDelim)
43+
if err != nil {
44+
return fmt.Errorf("invalid headers: %w", err)
45+
}
46+
// Note: CoAP headers would be mapped to options in the protocol implementation
47+
3948
sendOnce := func() {
4049
var body []byte
4150
var ct string
4251

43-
b, err := testpayload.Interpolate(sendPayload)
52+
b, err := testpayload.InterpolateWithDelimiters(sendPayload, openDelim, closeDelim)
4453
if err != nil {
4554
fmt.Fprintf(os.Stderr, "Failed to interpolate payload: %v\n", err)
4655
return
@@ -126,6 +135,8 @@ func sendCommand() *cobra.Command {
126135
toolutil.AddPayloadFlags(cmd, &sendPayload, "{}", &sendMIME, toolutil.CTJSON)
127136
toolutil.AddIntervalFlag(cmd, &sendInterval, "5s")
128137
cmd.Flags().StringVar(&sendProto, "proto", "udp", "CoAP transport protocol: udp or tcp")
138+
toolutil.AddHeadersFlag(cmd, &headers)
139+
toolutil.AddTemplateDelimiterFlags(cmd, &openDelim, &closeDelim)
129140

130141
return cmd
131142
}

‎httptool/send.go‎

Lines changed: 20 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,15 @@ import (
1212

1313
func sendCommand() *cobra.Command {
1414
var (
15-
address string
16-
method string
17-
path string
18-
payload string
19-
interval string
20-
mime string
15+
address string
16+
method string
17+
path string
18+
payload string
19+
interval string
20+
mime string
21+
headers []string
22+
openDelim string
23+
closeDelim string
2124
)
2225

2326
cmd := &cobra.Command{
@@ -33,8 +36,13 @@ func sendCommand() *cobra.Command {
3336
toolutil.PrintKeyValue("URL", url)
3437
toolutil.PrintKeyValue("Interval", interval)
3538

39+
headerMap, err := toolutil.ParseHeadersWithDelimiters(headers, openDelim, closeDelim)
40+
if err != nil {
41+
return fmt.Errorf("invalid headers: %w", err)
42+
}
43+
3644
sendRequest := func() {
37-
reqBody, contentType, err := toolutil.BuildPayload(payload, mime)
45+
reqBody, contentType, err := toolutil.BuildPayloadWithDelimiters(payload, mime, openDelim, closeDelim)
3846
if err != nil {
3947
fmt.Fprintln(os.Stderr, err)
4048
return
@@ -52,6 +60,9 @@ func sendCommand() *cobra.Command {
5260
if contentType != "" {
5361
r.Header.Set("Content-Type", contentType)
5462
}
63+
for k, v := range headerMap {
64+
r.Header.Set(k, v)
65+
}
5566
if len(reqBody) > 0 {
5667
r.SetBody(reqBody)
5768
}
@@ -77,6 +88,8 @@ func sendCommand() *cobra.Command {
7788
toolutil.AddPathFlag(cmd, &path, "/event", "HTTP request path")
7889
toolutil.AddPayloadFlags(cmd, &payload, "{}", &mime, toolutil.CTJSON)
7990
toolutil.AddIntervalFlag(cmd, &interval, "5s")
91+
toolutil.AddHeadersFlag(cmd, &headers)
92+
toolutil.AddTemplateDelimiterFlags(cmd, &openDelim, &closeDelim)
8093

8194
return cmd
8295
}

‎kafkatool/send.go‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,9 @@ func sendCommand() *cobra.Command {
2020
sendPayload string
2121
sendMIME string
2222
sendInterval string
23+
headers []string
24+
openDelim string
25+
closeDelim string
2326
)
2427

2528
cmd := &cobra.Command{
@@ -44,17 +47,27 @@ func sendCommand() *cobra.Command {
4447
ticker := time.NewTicker(dur)
4548
defer ticker.Stop()
4649

50+
headerMap, err := toolutil.ParseHeadersWithDelimiters(headers, openDelim, closeDelim)
51+
if err != nil {
52+
return fmt.Errorf("invalid headers: %w", err)
53+
}
54+
4755
logger := toolutil.Logger()
4856
logger.Info("Producing to Kafka", "brokers", sendBrokers, "topic", sendTopic, "interval", sendInterval)
4957

5058
for range ticker.C {
51-
body, _, err := toolutil.BuildPayload(sendPayload, sendMIME)
59+
body, _, err := toolutil.BuildPayloadWithDelimiters(sendPayload, sendMIME, openDelim, closeDelim)
5260
if err != nil {
5361
logger.Error("Failed to build payload", "error", err)
5462
continue
5563
}
64+
msg := kafka.Message{Value: body}
65+
for k, v := range headerMap {
66+
msg.Headers = append(msg.Headers, kafka.Header{Key: k, Value: []byte(v)})
67+
}
68+
5669
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
57-
err = w.WriteMessages(ctx, kafka.Message{Value: body})
70+
err = w.WriteMessages(ctx, msg)
5871
cancel()
5972
if err != nil {
6073
logger.Error("Failed to send message", "error", err)
@@ -70,6 +83,8 @@ func sendCommand() *cobra.Command {
7083
cmd.Flags().StringVar(&sendTopic, "topic", "test", "Kafka topic")
7184
toolutil.AddPayloadFlags(cmd, &sendPayload, "Hello, Kafka!", &sendMIME, toolutil.CTText)
7285
toolutil.AddIntervalFlag(cmd, &sendInterval, "5s")
86+
toolutil.AddHeadersFlag(cmd, &headers)
87+
toolutil.AddTemplateDelimiterFlags(cmd, &openDelim, &closeDelim)
7388

7489
return cmd
7590
}

‎mqtttool/send.go‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,9 @@ func sendCommand() *cobra.Command {
2727
sendQoS int
2828
sendRetain bool
2929
sendClientID string
30+
headers []string
31+
openDelim string
32+
closeDelim string
3033
)
3134

3235
cmd := &cobra.Command{
@@ -56,8 +59,14 @@ func sendCommand() *cobra.Command {
5659
toolutil.PrintKeyValue("QoS", sendQoS)
5760
toolutil.PrintKeyValue("Interval", sendInterval)
5861

62+
_, errHeaders := toolutil.ParseHeadersWithDelimiters(headers, openDelim, closeDelim)
63+
if errHeaders != nil {
64+
return fmt.Errorf("invalid headers: %w", errHeaders)
65+
}
66+
// Note: MQTT v5 user properties can be set from headers
67+
5968
publish := func() error {
60-
body, _, err := toolutil.BuildPayload(sendPayload, sendMIME)
69+
body, _, err := toolutil.BuildPayloadWithDelimiters(sendPayload, sendMIME, openDelim, closeDelim)
6170
if err != nil {
6271
toolutil.PrintError("Payload build error: %v", err)
6372
return err
@@ -83,6 +92,8 @@ func sendCommand() *cobra.Command {
8392
cmd.Flags().StringVar(&sendClientID, "clientid", "", "Client ID (auto if empty)")
8493
toolutil.AddPayloadFlags(cmd, &sendPayload, "{nowtime}", &sendMIME, toolutil.CTText)
8594
toolutil.AddIntervalFlag(cmd, &sendInterval, "5s")
95+
toolutil.AddHeadersFlag(cmd, &headers)
96+
toolutil.AddTemplateDelimiterFlags(cmd, &openDelim, &closeDelim)
8697

8798
return cmd
8899
}

‎natstool/send.go‎

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,9 @@ func sendCommand() *cobra.Command {
1717
sendMIME string
1818
sendInterval string
1919
sendStream string
20+
headers []string
21+
openDelim string
22+
closeDelim string
2023
)
2124

2225
cmd := &cobra.Command{
@@ -33,6 +36,11 @@ func sendCommand() *cobra.Command {
3336
defer nc.Close()
3437

3538
var js nats.JetStreamContext
39+
headerMap, err := toolutil.ParseHeadersWithDelimiters(headers, openDelim, closeDelim)
40+
if err != nil {
41+
return fmt.Errorf("invalid headers: %w", err)
42+
}
43+
3644
if sendStream != "" {
3745
if js, err = nc.JetStream(); err != nil {
3846
return fmt.Errorf("JetStream context error: %w", err)
@@ -48,20 +56,28 @@ func sendCommand() *cobra.Command {
4856
}
4957

5058
publish := func() error {
51-
body, _, err := toolutil.BuildPayload(sendPayload, sendMIME)
59+
body, _, err := toolutil.BuildPayloadWithDelimiters(sendPayload, sendMIME, openDelim, closeDelim)
5260
if err != nil {
5361
toolutil.PrintError("Payload build error: %v", err)
5462
return err
5563
}
64+
65+
// Build NATS message with headers
66+
msg := nats.NewMsg(sendSubject)
67+
msg.Data = body
68+
for k, v := range headerMap {
69+
msg.Header.Add(k, v)
70+
}
71+
5672
if sendStream != "" {
57-
ack, err := js.Publish(sendSubject, body)
73+
ack, err := js.PublishMsg(msg)
5874
if err != nil {
5975
toolutil.PrintError("JetStream publish error: %v", err)
6076
return err
6177
}
6278
toolutil.PrintInfo("Published to JetStream, sequence: %d", ack.Sequence)
6379
} else {
64-
if err := nc.Publish(sendSubject, body); err != nil {
80+
if err := nc.PublishMsg(msg); err != nil {
6581
toolutil.PrintError("Publish error: %v", err)
6682
return err
6783
}
@@ -79,6 +95,8 @@ func sendCommand() *cobra.Command {
7995
toolutil.AddPayloadFlags(cmd, &sendPayload, "{nowtime}", &sendMIME, toolutil.CTText)
8096
toolutil.AddIntervalFlag(cmd, &sendInterval, "5s")
8197
cmd.Flags().StringVar(&sendStream, "stream", "", "JetStream stream name (if set, uses JetStream)")
98+
toolutil.AddHeadersFlag(cmd, &headers)
99+
toolutil.AddTemplateDelimiterFlags(cmd, &openDelim, &closeDelim)
82100

83101
return cmd
84102
}

‎pkg/testpayload/testpayload.go‎

Lines changed: 43 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import (
44
"encoding/json"
55
"fmt"
66
"math/rand"
7+
"os"
78
"strings"
89
"sync"
910
"time"
@@ -83,6 +84,12 @@ func GenerateCounter() int {
8384
}
8485

8586
func Interpolate(str string) ([]byte, error) {
87+
return InterpolateWithDelimiters(str, "{{", "}}")
88+
}
89+
90+
// InterpolateWithDelimiters performs template variable interpolation with custom delimiters
91+
// Supports placeholders: json, cbor, sentiment, sentence, datetime, nowtime, counter, file:/path
92+
func InterpolateWithDelimiters(str string, openDelim string, closeDelim string) ([]byte, error) {
8693
placeholders := map[string]TestPayloadType{
8794
"json": TestPayloadJSON,
8895
"cbor": TestPayloadCBOR,
@@ -95,7 +102,7 @@ func Interpolate(str string) ([]byte, error) {
95102

96103
result := str
97104
for key, typ := range placeholders {
98-
ph := "{" + key + "}"
105+
ph := openDelim + key + closeDelim
99106

100107
if str == ph {
101108
// If the entire string is just the placeholder, return the generated value directly
@@ -110,6 +117,41 @@ func Interpolate(str string) ([]byte, error) {
110117
result = strings.ReplaceAll(result, ph, string(val))
111118
}
112119
}
120+
121+
// Handle file:// placeholder
122+
filePrefix := openDelim + "file:"
123+
fileSuffix := closeDelim
124+
if strings.Contains(result, filePrefix) {
125+
for {
126+
startIdx := strings.Index(result, filePrefix)
127+
if startIdx == -1 {
128+
break
129+
}
130+
endIdx := strings.Index(result[startIdx:], fileSuffix)
131+
if endIdx == -1 {
132+
return nil, fmt.Errorf("unclosed file placeholder at position %d", startIdx)
133+
}
134+
endIdx += startIdx
135+
136+
// Extract file path
137+
filePath := result[startIdx+len(filePrefix) : endIdx]
138+
if filePath == "" {
139+
return nil, fmt.Errorf("empty file path in placeholder at position %d", startIdx)
140+
}
141+
142+
// Read file content
143+
// #nosec G304 -- reading file for test payload generation
144+
content, err := os.ReadFile(filePath)
145+
if err != nil {
146+
return nil, fmt.Errorf("failed to read file %s: %w", filePath, err)
147+
}
148+
149+
// Replace placeholder with file content
150+
placeholder := result[startIdx : endIdx+len(fileSuffix)]
151+
result = strings.Replace(result, placeholder, string(content), 1)
152+
}
153+
}
154+
113155
return []byte(result), nil
114156
}
115157

0 commit comments

Comments
 (0)