Skip to content

Commit 692b563

Browse files
Added unit tests for kafkaServer implementation (#357)
* Cleaned server.go * Removed redundant goroutine * Changes to initializing order to fix resource cleanup * Unit tests for kafka server * Cleaned up tests * Cleaned up warnings in kafka unit tests * Modified kafka unit tests to mock and test more meaningful parts of KafkaServer.go * Simplified unit tests * Removed redundant nil check * Pulled complexity out of kafka tests * Renamed to setupConsumer + added unit tests in CI * Pulls out harness code from tests into a helper and sufficiently comments it * Added field in lkc to track error counts. Updated error handling test to use this field * Better commenting of runConsumeLoop --------- Co-authored-by: RSYashwanth <[email protected]>
1 parent 72b4f14 commit 692b563

3 files changed

Lines changed: 368 additions & 17 deletions

File tree

.github/workflows/ci.yml

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,10 @@ jobs:
6262
working-directory: go/common
6363
run: go test -v
6464
timeout-minutes: 5
65+
- name: Test Go Worker
66+
working-directory: go/worker
67+
run: go test -v ./...
68+
timeout-minutes: 5
6569
- name: Test Python (SOCK)
6670
run: sudo env "PATH=$PATH" ./scripts/test.py --worker_type=sock
6771
timeout-minutes: 20

go/worker/event/kafkaServer.go

Lines changed: 33 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -24,21 +24,37 @@ type KafkaClient interface {
2424
Close()
2525
}
2626

27+
// LambdaInvoker abstracts the lambda invocation layer for testability
28+
type LambdaInvoker interface {
29+
Invoke(lambdaName string, w http.ResponseWriter, r *http.Request)
30+
}
31+
32+
// lambdaMgrInvoker wraps *lambda.LambdaMgr to implement LambdaInvoker
33+
type lambdaMgrInvoker struct {
34+
mgr *lambda.LambdaMgr
35+
}
36+
37+
func (i *lambdaMgrInvoker) Invoke(lambdaName string, w http.ResponseWriter, r *http.Request) {
38+
f := i.mgr.Get(lambdaName)
39+
f.Invoke(w, r)
40+
}
41+
2742
// LambdaKafkaConsumer manages Kafka consumption for a specific lambda function
2843
type LambdaKafkaConsumer struct {
29-
consumerName string // Unique name for this consumer
30-
lambdaName string // lambda function name
31-
kafkaTrigger *common.KafkaTrigger
32-
client KafkaClient // kgo.client implements the KafkaClient interface
33-
lambdaManager *lambda.LambdaMgr // Reference to lambda manager for direct calls
34-
stopChan chan struct{} // Shutdown signal for this consumer
44+
consumerName string // Unique name for this consumer
45+
lambdaName string // lambda function name
46+
kafkaTrigger *common.KafkaTrigger
47+
client KafkaClient // kgo.client implements the KafkaClient interface
48+
invoker LambdaInvoker // Abstraction for lambda invocation
49+
stopChan chan struct{} // Shutdown signal for this consumer
3550
// When this channel is closed, the goroutine for the consumer exits
51+
errorCount int // Number of non-timeout Kafka client errors encountered
3652
}
3753

3854
// KafkaManager manages multiple lambda-specific Kafka consumers
3955
type KafkaManager struct {
4056
lambdaConsumers map[string]*LambdaKafkaConsumer // lambdaName -> consumer
41-
lambdaManager *lambda.LambdaMgr // Reference to lambda manager
57+
invoker LambdaInvoker // Abstraction for lambda invocation
4258
mu sync.Mutex // Protects lambdaConsumers map
4359
}
4460

@@ -75,20 +91,20 @@ func (km *KafkaManager) newLambdaKafkaConsumer(consumerName string, lambdaName s
7591
}
7692

7793
return &LambdaKafkaConsumer{
78-
consumerName: consumerName,
79-
lambdaName: lambdaName,
80-
kafkaTrigger: trigger,
81-
client: client,
82-
lambdaManager: km.lambdaManager,
83-
stopChan: make(chan struct{}),
94+
consumerName: consumerName,
95+
lambdaName: lambdaName,
96+
kafkaTrigger: trigger,
97+
client: client,
98+
invoker: km.invoker,
99+
stopChan: make(chan struct{}),
84100
}, nil
85101
}
86102

87103
// NewKafkaManager creates and configures a new Kafka manager
88104
func NewKafkaManager(lambdaManager *lambda.LambdaMgr) (*KafkaManager, error) {
89105
manager := &KafkaManager{
90106
lambdaConsumers: make(map[string]*LambdaKafkaConsumer),
91-
lambdaManager: lambdaManager,
107+
invoker: &lambdaMgrInvoker{mgr: lambdaManager},
92108
}
93109

94110
slog.Info("Kafka manager initialized")
@@ -129,6 +145,7 @@ func (lkc *LambdaKafkaConsumer) consumeLoop() {
129145
continue
130146
}
131147

148+
lkc.errorCount++
132149
// TODO: Surface Kafka consumer errors to lambda developers by invoking an error
133150
// handler lambda function. Could allow lambdas to specify an onError callback in
134151
// ol.yaml that gets invoked with error details.
@@ -185,9 +202,8 @@ func (lkc *LambdaKafkaConsumer) processMessage(record *kgo.Record) {
185202
// for kafka triggered lambda invocations.
186203
w := httptest.NewRecorder()
187204

188-
// Get lambda function and invoke directly
189-
lambdaFunc := lkc.lambdaManager.Get(lkc.lambdaName)
190-
lambdaFunc.Invoke(w, req)
205+
// Invoke the lambda function directly
206+
lkc.invoker.Invoke(lkc.lambdaName, w, req)
191207

192208
// Log the result
193209
slog.Info("Kafka message processed via direct invocation",

0 commit comments

Comments
 (0)