Skip to content

Commit b095679

Browse files
committed
improve tests
1 parent 298d515 commit b095679

File tree

2 files changed

+166
-84
lines changed

2 files changed

+166
-84
lines changed

pkg/decision/cmab_client.go

+18-5
Original file line numberDiff line numberDiff line change
@@ -127,11 +127,17 @@ func NewDefaultCmabClient(options CmabClientOptions) *DefaultCmabClient {
127127

128128
// FetchDecision fetches a decision from the CMAB API
129129
func (c *DefaultCmabClient) FetchDecision(
130+
ctx context.Context,
130131
ruleID string,
131132
userID string,
132133
attributes map[string]interface{},
133134
cmabUUID string,
134135
) (string, error) {
136+
// If no context is provided, create a background context
137+
if ctx == nil {
138+
ctx = context.Background()
139+
}
140+
135141
// Create the URL
136142
url := fmt.Sprintf(CMABPredictionEndpoint, ruleID)
137143

@@ -163,16 +169,18 @@ func (c *DefaultCmabClient) FetchDecision(
163169
return "", fmt.Errorf("failed to marshal CMAB request: %w", err)
164170
}
165171

166-
// Create context for cancellation
167-
ctx := context.Background()
168-
169172
// If no retry config, just do a single fetch
170173
if c.retryConfig == nil {
171174
return c.doFetch(ctx, url, bodyBytes)
172175
}
173176

174177
// Retry sending request with exponential backoff
175178
for i := 0; i <= c.retryConfig.MaxRetries; i++ {
179+
// Check if context is done
180+
if ctx.Err() != nil {
181+
return "", fmt.Errorf("context canceled or timed out: %w", ctx.Err())
182+
}
183+
176184
// Make the request
177185
result, err := c.doFetch(ctx, url, bodyBytes)
178186
if err == nil {
@@ -194,8 +202,13 @@ func (c *DefaultCmabClient) FetchDecision(
194202
c.logger.Debug(fmt.Sprintf("CMAB request retry %d/%d, backing off for %v",
195203
i+1, c.retryConfig.MaxRetries, backoffDuration))
196204

197-
// Wait for backoff duration
198-
time.Sleep(backoffDuration)
205+
// Wait for backoff duration with context awareness
206+
select {
207+
case <-ctx.Done():
208+
return "", fmt.Errorf("context canceled or timed out during backoff: %w", ctx.Err())
209+
case <-time.After(backoffDuration):
210+
// Continue with retry
211+
}
199212

200213
c.logger.Warning(fmt.Sprintf("CMAB API request failed (attempt %d/%d): %v",
201214
i+1, c.retryConfig.MaxRetries, err))

pkg/decision/cmab_client_test.go

+148-79
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
package decision
1919

2020
import (
21+
"context"
2122
"encoding/json"
2223
"fmt"
2324
"io"
@@ -75,20 +76,36 @@ func TestDefaultCmabClient_FetchDecision(t *testing.T) {
7576
assert.Equal(t, "rule456", instance.ExperimentID)
7677
assert.Equal(t, "test-uuid", instance.CmabUUID)
7778

78-
// Verify attributes
79-
assert.Len(t, instance.Attributes, 2)
80-
// Attributes order is not guaranteed, so we need to check both
79+
// Verify attributes - check for various types
80+
assert.Len(t, instance.Attributes, 5)
81+
82+
// Create a map for easier attribute checking
83+
attrMap := make(map[string]CMABAttribute)
8184
for _, attr := range instance.Attributes {
82-
if attr.ID == "browser" {
83-
assert.Equal(t, "chrome", attr.Value)
84-
} else if attr.ID == "isMobile" {
85-
assert.Equal(t, true, attr.Value)
86-
} else {
87-
t.Errorf("Unexpected attribute ID: %s", attr.ID)
88-
}
85+
attrMap[attr.ID] = attr
8986
assert.Equal(t, "custom_attribute", attr.Type)
9087
}
9188

89+
// Check string attribute
90+
assert.Contains(t, attrMap, "string_attr")
91+
assert.Equal(t, "string value", attrMap["string_attr"].Value)
92+
93+
// Check int attribute
94+
assert.Contains(t, attrMap, "int_attr")
95+
assert.Equal(t, float64(42), attrMap["int_attr"].Value) // JSON numbers are float64
96+
97+
// Check float attribute
98+
assert.Contains(t, attrMap, "float_attr")
99+
assert.Equal(t, 3.14, attrMap["float_attr"].Value)
100+
101+
// Check bool attribute
102+
assert.Contains(t, attrMap, "bool_attr")
103+
assert.Equal(t, true, attrMap["bool_attr"].Value)
104+
105+
// Check null attribute
106+
assert.Contains(t, attrMap, "null_attr")
107+
assert.Nil(t, attrMap["null_attr"].Value)
108+
92109
// Return response
93110
w.Header().Set("Content-Type", "application/json")
94111
w.WriteHeader(http.StatusOK)
@@ -115,12 +132,19 @@ func TestDefaultCmabClient_FetchDecision(t *testing.T) {
115132
CMABPredictionEndpoint = server.URL + "/%s"
116133
defer func() { CMABPredictionEndpoint = originalEndpoint }()
117134

118-
// Test fetch decision
135+
// Test with various attribute types
119136
attributes := map[string]interface{}{
120-
"browser": "chrome",
121-
"isMobile": true,
137+
"string_attr": "string value",
138+
"int_attr": 42,
139+
"float_attr": 3.14,
140+
"bool_attr": true,
141+
"null_attr": nil,
122142
}
123-
variationID, err := client.FetchDecision("rule456", "user123", attributes, "test-uuid")
143+
144+
// Create a context for the request
145+
ctx := context.Background()
146+
147+
variationID, err := client.FetchDecision(ctx, "rule456", "user123", attributes, "test-uuid")
124148

125149
// Verify results
126150
assert.NoError(t, err)
@@ -199,7 +223,7 @@ func TestDefaultCmabClient_FetchDecision_WithRetry(t *testing.T) {
199223
}
200224

201225
startTime := time.Now()
202-
variationID, err := client.FetchDecision("rule456", "user123", attributes, "test-uuid")
226+
variationID, err := client.FetchDecision(context.Background(), "rule456", "user123", attributes, "test-uuid")
203227
duration := time.Since(startTime)
204228

205229
// Verify results
@@ -246,7 +270,7 @@ func TestDefaultCmabClient_FetchDecision_ExhaustedRetries(t *testing.T) {
246270
"isMobile": true,
247271
}
248272

249-
variationID, err := client.FetchDecision("rule456", "user123", attributes, "test-uuid")
273+
variationID, err := client.FetchDecision(context.Background(), "rule456", "user123", attributes, "test-uuid")
250274

251275
// Verify results
252276
assert.Error(t, err)
@@ -285,7 +309,7 @@ func TestDefaultCmabClient_FetchDecision_NoRetryConfig(t *testing.T) {
285309
"browser": "chrome",
286310
}
287311

288-
_, err := client.FetchDecision("rule456", "user123", attributes, "test-uuid")
312+
_, err := client.FetchDecision(context.Background(), "rule456", "user123", attributes, "test-uuid")
289313

290314
// Verify results
291315
assert.Error(t, err)
@@ -348,7 +372,7 @@ func TestDefaultCmabClient_FetchDecision_InvalidResponse(t *testing.T) {
348372
"browser": "chrome",
349373
}
350374

351-
_, err := client.FetchDecision("rule456", "user123", attributes, "test-uuid")
375+
_, err := client.FetchDecision(context.Background(), "rule456", "user123", attributes, "test-uuid")
352376

353377
// Verify results
354378
assert.Error(t, err)
@@ -392,7 +416,7 @@ func TestDefaultCmabClient_FetchDecision_NetworkErrors(t *testing.T) {
392416
"browser": "chrome",
393417
}
394418

395-
_, err := client.FetchDecision("rule456", "user123", attributes, "test-uuid")
419+
_, err := client.FetchDecision(context.Background(), "rule456", "user123", attributes, "test-uuid")
396420

397421
// Verify results
398422
assert.Error(t, err)
@@ -452,7 +476,7 @@ func TestDefaultCmabClient_ExponentialBackoff(t *testing.T) {
452476
"browser": "chrome",
453477
}
454478

455-
variationID, err := client.FetchDecision("rule456", "user123", attributes, "test-uuid")
479+
variationID, err := client.FetchDecision(context.Background(), "rule456", "user123", attributes, "test-uuid")
456480

457481
// Verify results
458482
require.NoError(t, err)
@@ -478,64 +502,6 @@ func TestDefaultCmabClient_ExponentialBackoff(t *testing.T) {
478502
}
479503
}
480504

481-
func TestDefaultCmabClient_RequestValidation(t *testing.T) {
482-
// Setup test server that validates the request format
483-
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
484-
// Read and validate the request body
485-
body, err := io.ReadAll(r.Body)
486-
require.NoError(t, err)
487-
488-
// Check that the body is valid JSON
489-
var requestBody map[string]interface{}
490-
err = json.Unmarshal(body, &requestBody)
491-
require.NoError(t, err)
492-
493-
// Check that the required fields are present
494-
instances, ok := requestBody["instances"].([]interface{})
495-
require.True(t, ok, "Request should have 'instances' array")
496-
require.Len(t, instances, 1, "Request should have exactly one instance")
497-
498-
instance := instances[0].(map[string]interface{})
499-
require.Contains(t, instance, "visitorId", "Instance should have visitorId")
500-
require.Contains(t, instance, "experimentId", "Instance should have experimentId")
501-
require.Contains(t, instance, "attributes", "Instance should have attributes")
502-
require.Contains(t, instance, "cmabUUID", "Instance should have cmabUUID")
503-
504-
// Return success response
505-
w.Header().Set("Content-Type", "application/json")
506-
w.WriteHeader(http.StatusOK)
507-
fmt.Fprintf(w, `{"predictions":[{"variation_id":"var123"}]}`)
508-
}))
509-
defer server.Close()
510-
511-
// Create client with custom endpoint
512-
client := NewDefaultCmabClient(CmabClientOptions{
513-
HTTPClient: &http.Client{
514-
Timeout: 5 * time.Second,
515-
},
516-
})
517-
518-
// Override the endpoint for testing
519-
originalEndpoint := CMABPredictionEndpoint
520-
CMABPredictionEndpoint = server.URL + "/%s"
521-
defer func() { CMABPredictionEndpoint = originalEndpoint }()
522-
523-
// Test with various attribute types
524-
attributes := map[string]interface{}{
525-
"string_attr": "string value",
526-
"int_attr": 42,
527-
"float_attr": 3.14,
528-
"bool_attr": true,
529-
"null_attr": nil,
530-
}
531-
532-
variationID, err := client.FetchDecision("rule456", "user123", attributes, "test-uuid")
533-
534-
// Verify results
535-
assert.NoError(t, err)
536-
assert.Equal(t, "var123", variationID)
537-
}
538-
539505
func TestNewDefaultCmabClient_DefaultValues(t *testing.T) {
540506
// Test with empty options
541507
client := NewDefaultCmabClient(CmabClientOptions{})
@@ -598,7 +564,7 @@ func TestDefaultCmabClient_LoggingBehavior(t *testing.T) {
598564
"browser": "chrome",
599565
}
600566

601-
_, err := client.FetchDecision("rule456", "user123", attributes, "test-uuid")
567+
_, err := client.FetchDecision(context.Background(), "rule456", "user123", attributes, "test-uuid")
602568
assert.NoError(t, err)
603569

604570
// Verify log messages
@@ -619,3 +585,106 @@ func TestDefaultCmabClient_LoggingBehavior(t *testing.T) {
619585
assert.True(t, foundRetryWarning, "Expected warning log about API request failure")
620586
assert.True(t, foundBackoffDebug, "Expected debug log about retry backoff")
621587
}
588+
589+
func TestDefaultCmabClient_NonSuccessStatusCode(t *testing.T) {
590+
// Setup test server that returns different non-2xx status codes
591+
testCases := []struct {
592+
name string
593+
statusCode int
594+
statusText string
595+
}{
596+
{"BadRequest", http.StatusBadRequest, "Bad Request"},
597+
{"Unauthorized", http.StatusUnauthorized, "Unauthorized"},
598+
{"Forbidden", http.StatusForbidden, "Forbidden"},
599+
{"NotFound", http.StatusNotFound, "Not Found"},
600+
{"InternalServerError", http.StatusInternalServerError, "Internal Server Error"},
601+
{"ServiceUnavailable", http.StatusServiceUnavailable, "Service Unavailable"},
602+
}
603+
604+
for _, tc := range testCases {
605+
t.Run(tc.name, func(t *testing.T) {
606+
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
607+
w.WriteHeader(tc.statusCode)
608+
w.Write([]byte(tc.statusText))
609+
}))
610+
defer server.Close()
611+
612+
// Create client with custom endpoint and no retries
613+
client := NewDefaultCmabClient(CmabClientOptions{
614+
HTTPClient: &http.Client{
615+
Timeout: 5 * time.Second,
616+
},
617+
// No retry config to simplify the test
618+
})
619+
620+
// Override the endpoint for testing
621+
originalEndpoint := CMABPredictionEndpoint
622+
CMABPredictionEndpoint = server.URL + "/%s"
623+
defer func() { CMABPredictionEndpoint = originalEndpoint }()
624+
625+
// Test fetch decision
626+
attributes := map[string]interface{}{
627+
"browser": "chrome",
628+
}
629+
630+
// Create a context for the request
631+
ctx := context.Background()
632+
633+
variationID, err := client.FetchDecision(ctx, "rule456", "user123", attributes, "test-uuid")
634+
635+
// Verify results
636+
assert.Error(t, err, "Expected error for non-success status code")
637+
assert.Equal(t, "", variationID, "Expected empty variation ID for error response")
638+
assert.Contains(t, err.Error(), "non-success status code")
639+
assert.Contains(t, err.Error(), fmt.Sprintf("%d", tc.statusCode))
640+
})
641+
}
642+
}
643+
644+
func TestDefaultCmabClient_FetchDecision_ContextCancellation(t *testing.T) {
645+
// Setup test server that delays response
646+
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
647+
// Sleep to simulate a slow response
648+
time.Sleep(500 * time.Millisecond)
649+
650+
w.Header().Set("Content-Type", "application/json")
651+
w.WriteHeader(http.StatusOK)
652+
response := CMABResponse{
653+
Predictions: []CMABPrediction{
654+
{
655+
VariationID: "var123",
656+
},
657+
},
658+
}
659+
json.NewEncoder(w).Encode(response)
660+
}))
661+
defer server.Close()
662+
663+
// Create client with custom endpoint
664+
client := NewDefaultCmabClient(CmabClientOptions{
665+
HTTPClient: &http.Client{
666+
Timeout: 5 * time.Second,
667+
},
668+
})
669+
670+
// Override the endpoint for testing
671+
originalEndpoint := CMABPredictionEndpoint
672+
CMABPredictionEndpoint = server.URL + "/%s"
673+
defer func() { CMABPredictionEndpoint = originalEndpoint }()
674+
675+
// Create a context with a short timeout
676+
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
677+
defer cancel()
678+
679+
// Test fetch decision with a context that will time out
680+
attributes := map[string]interface{}{
681+
"browser": "chrome",
682+
}
683+
684+
_, err := client.FetchDecision(ctx, "rule456", "user123", attributes, "test-uuid")
685+
686+
// Verify that we got a context deadline exceeded error
687+
assert.Error(t, err)
688+
assert.Contains(t, err.Error(), "context")
689+
assert.Contains(t, err.Error(), "deadline exceeded")
690+
}

0 commit comments

Comments
 (0)