Skip to content

Commit bbc57dc

Browse files
vikin91vikin91
committed
Address review: logging & test coverage
1 parent aba2b54 commit bbc57dc

File tree

2 files changed

+63
-9
lines changed

2 files changed

+63
-9
lines changed

‎central/sensor/service/pipeline/virtualmachineindex/pipeline.go‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -130,7 +130,7 @@ func (p *pipelineImpl) Run(ctx context.Context, clusterID string, msg *central.M
130130
if p.rateLimiter == nil {
131131
log.Warnf("No rate limiter found for %s. Dropping VM index report %s from cluster %s", rateLimiterWorkload, index.GetId(), clusterID)
132132
if conn != nil && conn.HasCapability(centralsensor.SensorACKSupport) {
133-
sendVMIndexReportResponse(ctx, index.GetId(), central.SensorACK_NACK, "nil rate limiter", injector)
133+
sendVMIndexReportResponse(ctx, clusterID, index.GetId(), central.SensorACK_NACK, "nil rate limiter", injector)
134134
}
135135
return nil // Don't return error - would cause pipeline retry
136136
}
@@ -139,7 +139,7 @@ func (p *pipelineImpl) Run(ctx context.Context, clusterID string, msg *central.M
139139
if !allowed {
140140
log.Infof("Dropping VM index report %s from cluster %s: %s", index.GetId(), clusterID, reason)
141141
if conn != nil && conn.HasCapability(centralsensor.SensorACKSupport) {
142-
sendVMIndexReportResponse(ctx, index.GetId(), central.SensorACK_NACK, reason, injector)
142+
sendVMIndexReportResponse(ctx, clusterID, index.GetId(), central.SensorACK_NACK, reason, injector)
143143
}
144144
return nil // Don't return error - would cause pipeline retry
145145
}
@@ -169,15 +169,15 @@ func (p *pipelineImpl) Run(ctx context.Context, clusterID string, msg *central.M
169169

170170
// Send ACK to Sensor if Sensor supports it
171171
if conn != nil && conn.HasCapability(centralsensor.SensorACKSupport) {
172-
sendVMIndexReportResponse(ctx, index.GetId(), central.SensorACK_ACK, "", injector)
172+
sendVMIndexReportResponse(ctx, clusterID, index.GetId(), central.SensorACK_ACK, "", injector)
173173
}
174174
return nil
175175
}
176176

177177
// sendVMIndexReportResponse sends an ACK or NACK for a VM index report.
178-
func sendVMIndexReportResponse(ctx context.Context, vmID string, action central.SensorACK_Action, reason string, injector common.MessageInjector) {
178+
func sendVMIndexReportResponse(ctx context.Context, clusterID, vmID string, action central.SensorACK_Action, reason string, injector common.MessageInjector) {
179179
if injector == nil {
180-
log.Debugf("Cannot send %s to Sensor - no injector", action.String())
180+
log.Debugf("Cannot send %s to Sensor for cluster %s - no injector", action.String(), clusterID)
181181
return
182182
}
183183
msg := &central.MsgToSensor{
@@ -191,8 +191,8 @@ func sendVMIndexReportResponse(ctx context.Context, vmID string, action central.
191191
},
192192
}
193193
if err := injector.InjectMessage(ctx, msg); err != nil {
194-
log.Warnf("Failed sending VM index report %s for %s: %v", action.String(), vmID, err)
194+
log.Warnf("Failed sending VM index report %s for VM %s in cluster %s: %v", action.String(), vmID, clusterID, err)
195195
} else {
196-
log.Debugf("Sent VM index report %s for %s (reason=%q)", action.String(), vmID, reason)
196+
log.Debugf("Sent VM index report %s for VM %s in cluster %s (reason=%q)", action.String(), vmID, clusterID, reason)
197197
}
198198
}

‎central/sensor/service/pipeline/virtualmachineindex/pipeline_test.go‎

Lines changed: 56 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -388,7 +388,7 @@ func TestPipelineRun_RateLimitDisabled(t *testing.T) {
388388
msg := createVMIndexMessage(vmID, central.ResourceAction_SYNC_RESOURCE)
389389

390390
// Should process all 100 requests without rate limiting
391-
for i := 0; i < 100; i++ {
391+
for i := range 100 {
392392
enricher.EXPECT().
393393
EnrichVirtualMachineWithVulnerabilities(gomock.Any(), gomock.Any()).
394394
Return(nil)
@@ -448,7 +448,7 @@ func TestPipelineRun_RateLimitEnabled(t *testing.T) {
448448

449449
// Send 6 requests - the first 5 should be processed successfully,
450450
// the 6th should be rate-limited.
451-
for i := 0; i < 6; i++ {
451+
for i := range 6 {
452452
err := pipeline.Run(ctxWithConn, testClusterID, msg, injector)
453453
assert.NoError(t, err, "Run should not return an error even when rate-limited (request %d)", i+1)
454454
}
@@ -469,6 +469,60 @@ func TestPipelineRun_RateLimitEnabled(t *testing.T) {
469469
assert.Contains(t, acks[5].GetReason(), "rate limit exceeded")
470470
}
471471

472+
// TestPipelineRun_RateLimitEnabled_NoACKSupport tests that when the connection
473+
// does not support SensorACKSupport, rate limiting still applies but no ACK/NACK
474+
// messages are sent.
475+
func TestPipelineRun_RateLimitEnabled_NoACKSupport(t *testing.T) {
476+
t.Setenv(features.VirtualMachines.EnvVar(), "true")
477+
ctrl := gomock.NewController(t)
478+
defer ctrl.Finish()
479+
480+
vmDatastore := vmDatastoreMocks.NewMockDataStore(ctrl)
481+
enricher := vmEnricherMocks.NewMockVirtualMachineEnricher(ctrl)
482+
rateLimiter := mustNewLimiter(t, "test-no-ack", 5, 5) // 5 req/s, bucket capacity=5
483+
484+
// Recording injector to capture any ACK/NACK attempts - should remain empty.
485+
injector := &recordingInjector{}
486+
487+
// Mock connection WITHOUT SensorACKSupport capability
488+
mockConn := connMocks.NewMockSensorConnection(ctrl)
489+
mockConn.EXPECT().HasCapability(centralsensor.SensorACKSupport).Return(false).AnyTimes()
490+
491+
pipeline := &pipelineImpl{
492+
vmDatastore: vmDatastore,
493+
enricher: enricher,
494+
rateLimiter: rateLimiter,
495+
}
496+
497+
vmID := "vm-1"
498+
msg := createVMIndexMessage(vmID, central.ResourceAction_SYNC_RESOURCE)
499+
500+
// Build a context with the mocked connection that does NOT have SensorACKSupport
501+
ctxWithConn := connection.WithConnection(context.Background(), mockConn)
502+
503+
// Expect enrichment and datastore writes ONLY for the first 5 (non-rate-limited) requests.
504+
enricher.EXPECT().
505+
EnrichVirtualMachineWithVulnerabilities(gomock.Any(), gomock.Any()).
506+
Return(nil).
507+
Times(5)
508+
509+
vmDatastore.EXPECT().
510+
UpdateVirtualMachineScan(gomock.Any(), vmID, gomock.Any()).
511+
Return(nil).
512+
Times(5)
513+
514+
// Send 6 requests - the first 5 should be processed successfully,
515+
// the 6th should be rate-limited.
516+
for i := range 6 {
517+
err := pipeline.Run(ctxWithConn, testClusterID, msg, injector)
518+
assert.NoError(t, err, "Run should not return an error even when rate-limited (request %d)", i+1)
519+
}
520+
521+
// Verify NO ACK/NACK messages were sent (SensorACKSupport is not available)
522+
acks := injector.getSentACKs()
523+
assert.Empty(t, acks, "no ACK/NACKs should be sent when SensorACKSupport is not available")
524+
}
525+
472526
// TestPipelineRun_NilRateLimiter_WithACKSupport tests behavior when the rateLimiter is nil and ACKs are supported.
473527
// This covers the nil-limiter branch and verifies that:
474528
// 1. No enrichment/datastore calls occur

0 commit comments

Comments
 (0)