@@ -218,6 +218,46 @@ func TestStatsEngineAddRemoveContainers(t *testing.T) {
218218 validateIdleContainerMetrics (t , engine )
219219}
220220
221+ func TestDoRemoveContainerUnsafeServiceConnectStats (t * testing.T ) {
222+ engine := NewDockerStatsEngine (& cfg , nil , eventStream ("TestDoRemoveContainerUnsafeServiceConnectStats" ), nil , nil , nil )
223+ defer engine .removeAll ()
224+
225+ taskArn := "arn:aws:ecs:us-west-2:123456789012:task/test-task"
226+
227+ // Setup task with service connect stats
228+ engine .taskToServiceConnectStats [taskArn ] = & ServiceConnectStats {}
229+ engine .tasksToContainers [taskArn ] = make (map [string ]* StatsContainer )
230+
231+ // Add two containers to the task with proper initialization
232+ _ , cancel1 := context .WithCancel (context .Background ())
233+ container1 := & StatsContainer {
234+ containerMetadata : & ContainerMetadata {DockerID : "container1" },
235+ cancel : cancel1 ,
236+ }
237+ _ , cancel2 := context .WithCancel (context .Background ())
238+ container2 := & StatsContainer {
239+ containerMetadata : & ContainerMetadata {DockerID : "container2" },
240+ cancel : cancel2 ,
241+ }
242+
243+ engine .tasksToContainers [taskArn ]["container1" ] = container1
244+ engine .tasksToContainers [taskArn ]["container2" ] = container2
245+
246+ // Remove first container - service connect stats should remain
247+ engine .doRemoveContainerUnsafe (container1 , taskArn )
248+
249+ // Verify service connect stats still exist (task has remaining containers)
250+ _ , exists := engine .taskToServiceConnectStats [taskArn ]
251+ assert .True (t , exists , "Service connect stats should remain when task has running containers" )
252+
253+ // Remove second container - service connect stats should be cleaned up
254+ engine .doRemoveContainerUnsafe (container2 , taskArn )
255+
256+ // Verify service connect stats are cleaned up (no containers left in task)
257+ _ , exists = engine .taskToServiceConnectStats [taskArn ]
258+ assert .False (t , exists , "Service connect stats should be cleaned up when no containers remain in task" )
259+ }
260+
221261func TestStatsEngineMetadataInStatsSets (t * testing.T ) {
222262 mockCtrl := gomock .NewController (t )
223263 defer mockCtrl .Finish ()
0 commit comments