// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. // // Licensed under the Apache License, Version 2.0 (the "License"). You may not // use this file except in compliance with the License. A copy of the // License is located at // // http://aws.amazon.com/apache2.0/ // // or in the "license" file accompanying this file. This file is distributed // on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, // either express or implied. See the License for the specific language governing // permissions and limitations under the License. //go:build integration // +build integration package registrar import ( "github.com/aws/amazon-ssm-agent/agent/appconfig" "github.com/aws/amazon-ssm-agent/agent/ssm/authregister" "github.com/aws/amazon-ssm-agent/common/identity/availableidentities/ec2" "github.com/aws/amazon-ssm-agent/common/identity/availableidentities/ec2/mocks" "github.com/aws/aws-sdk-go/aws" "github.com/stretchr/testify/mock" "net/http" "sync" "testing" "time" "github.com/aws/amazon-ssm-agent/agent/mocks/log" "github.com/stretchr/testify/assert" ) func TestRetryableRegistrar_RegisterWithRetry_WhenIMDSAvailable_AndSSMUnavailable_Cancelled(t *testing.T) { // Arrange log := log.NewMockLog() awsConfig := &aws.Config{} awsConfig = awsConfig.WithMaxRetries(3). WithEndpoint("www.google.com:81"). // Endpoint is unreachable which causes timeout WithHTTPClient(&http.Client{Timeout: time.Second * 10}) // Decrease timeout from http default for test efficiency imdsClient := &mocks.IEC2MdsSdkClient{} imdsClient.On("GetMetadataWithContext", mock.Anything, mock.Anything). Return("SomeInstanceId", nil). Repeatability = 0 // TODO: GetMetadata is still called by the IIRRoleProvider Retrieve() method instead of GetMetadataWithContext imdsClient.On("GetMetadata", mock.Anything). Return("SomeInstanceId", nil). Repeatability = 0 imdsClient.On("RegionWithContext", mock.Anything). Return("SomeRegion", nil). Repeatability = 0 config := appconfig.SsmagentConfig{ Ssm: appconfig.SsmCfg{}, } authRegisterService := authregister.NewClientWithConfig(log, config, imdsClient, *awsConfig) ec2Identity := &ec2.Identity{ Log: log, Client: imdsClient, Config: &config, AuthRegisterService: authRegisterService, } registrar := &RetryableRegistrar{ log: log, identityRegistrar: ec2Identity, registrationAttemptedChan: make(chan struct{}, 1), stopRegistrarChan: make(chan struct{}, 1), timeAfterFunc: time.After, isRegistrarRunning: true, isRegistrarRunningLock: &sync.RWMutex{}, } // Act complete := make(chan struct{}) go func() { registrar.RegisterWithRetry() complete <- struct{}{} close(complete) }() <-time.After(2 * time.Second) // Act registrar.Stop() // Assert select { case <-complete: case <-time.After(1 * time.Second): assert.Fail(t, "Test did not complete in allotted time") } // Assert assert.False(t, registrar.getIsRegistrarRunning(), "Registrar still running after Stop() called") select { case <-registrar.registrationAttemptedChan: // Registration attempted chan closed successfully default: assert.Fail(t, "registration not marked as attempted") } }