diff --git a/.travis.yml b/.travis.yml
index eb8b7b9..9fc00f2 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -5,5 +5,7 @@ go:
 before_install:
   - go get github.com/mattn/goveralls
   - go get golang.org/x/tools/cmd/cover
+  - go get honnef.co/go/tools/cmd/staticcheck
 script:
+  - $HOME/gopath/bin/staticcheck ./...
   - $HOME/gopath/bin/goveralls -service=travis-ci
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 0c98ddf..f279a28 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -12,10 +12,14 @@ Make sure your contribution passes the following validations:
 
    `golint ./...`
 
-3. And new code must pass Go Vetting practices:
+3. New code must pass Go Vetting practices:
 
    `go vet ./...`
 
+4. And new code must pass [staticcheck](https://godoc.org/honnef.co/go/tools/cmd/staticcheck) checks:
+
+   `staticcheck ./...`
+
 I would like to keep this library simple, the proposed change must be a common use case.
 
 ## maintainer
diff --git a/README.md b/README.md
index f906aff..208b5b8 100644
--- a/README.md
+++ b/README.md
@@ -1,5 +1,7 @@
 # go-awsecs
 
+[![godoc reference](http://img.shields.io/badge/godoc-reference-blue.svg)](https://pkg.go.dev/github.com/Autodesk/go-awsecs)
+
 [![travis ci](https://api.travis-ci.org/Autodesk/go-awsecs.svg?branch=master)](https://travis-ci.org/Autodesk/go-awsecs)
 
 [![coverage status](https://coveralls.io/repos/github/Autodesk/go-awsecs/badge.svg?branch=master)](https://coveralls.io/github/Autodesk/go-awsecs?branch=master)
diff --git a/asg.go b/asg.go
index 5020e98..ba68b46 100644
--- a/asg.go
+++ b/asg.go
@@ -4,6 +4,7 @@ import (
 	"errors"
 	"fmt"
 	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/arn"
 	"github.com/aws/aws-sdk-go/service/autoscaling"
 	"github.com/aws/aws-sdk-go/service/autoscaling/autoscalingiface"
 	"github.com/aws/aws-sdk-go/service/ec2"
@@ -170,16 +171,40 @@ func drainingContainerInstanceIsDrained(ECSAPI ecs.ECS, clusterName, containerIn
 	if err != nil {
 		return err
 	}
-	for _, containerInstance := range output.ContainerInstances {
-		if *containerInstance.Status != "DRAINING" {
-			return backoff.Permanent(errors.New("the instance should be DRAINING but is not"))
+	return findDrainingContainerInstance(output, containerInstanceID)
+}
+
+func findDrainingContainerInstance(output *ecs.DescribeContainerInstancesOutput, containerInstanceID string) error {
+	if len(output.ContainerInstances) == 0 {
+		return ErrContainerInstanceNotFound
+	}
+	containerInstance := output.ContainerInstances[0]
+	containerInstanceArn := *containerInstance.ContainerInstanceArn
+	parsedArn, err := arn.Parse(containerInstanceArn)
+	if err != nil {
+		return err
+	}
+	return checkDrainingContainerInstance(containerInstance, parsedArn, containerInstanceID)
+}
+
+func checkDrainingContainerInstance(containerInstance *ecs.ContainerInstance, parsedArn arn.ARN, containerInstanceID string) error {
+	containerInstanceIDFound := strings.TrimPrefix(parsedArn.Resource, "container-instance/")
+	if containerInstanceIDFound == containerInstanceID {
+		if *containerInstance.Status != ecs.ContainerInstanceStatusDraining {
+			errorStringFormat := "the instance should be %s but is not"
+			errorString := fmt.Sprintf(errorStringFormat, ecs.ContainerInstanceStatusDraining)
+			permanentError := errors.New(errorString)
+			return backoff.Permanent(permanentError)
 		}
 		if *containerInstance.RunningTasksCount != 0 {
-			return errors.New("container instance still DRAINING")
+			errorStringFormat := "container instance still %s"
+			errorString := fmt.Sprintf(errorStringFormat, ecs.ContainerInstanceStatusDraining)
+			retryableError := errors.New(errorString)
+			return retryableError
 		}
 		return nil
 	}
-	return backoff.Permanent(errors.New("container instance not found"))
+	return ErrContainerInstanceNotFound
 }
 
 func drainAll(ASAPI autoscaling.AutoScaling, ECSAPI ecs.ECS, EC2API ec2.EC2, instances []ecsEC2Instance, asgName, clusterName string) error {
diff --git a/asg_test.go b/asg_test.go
index c74b772..34c91b4 100644
--- a/asg_test.go
+++ b/asg_test.go
@@ -1,8 +1,11 @@
 package awsecs
 
 import (
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/arn"
 	"github.com/aws/aws-sdk-go/service/autoscaling"
 	"github.com/aws/aws-sdk-go/service/autoscaling/autoscalingiface"
+	"github.com/aws/aws-sdk-go/service/ecs"
 	"testing"
 )
 
@@ -54,3 +57,76 @@ func TestFilterInstancesToReplace(t *testing.T) {
 		t.Errorf("unexpected")
 	}
 }
+
+func TestCheckDrainingContainerInstance(t *testing.T) {
+	type args struct {
+		containerInstance   *ecs.ContainerInstance
+		parsedArn           arn.ARN
+		containerInstanceID string
+	}
+	tests := []struct {
+		name    string
+		wantErr bool
+		args    args
+	}{
+		{
+			name:    "Found container instance ACTIVE",
+			wantErr: true,
+			args: args{
+				containerInstance: &ecs.ContainerInstance{
+					Status: aws.String(ecs.ContainerInstanceStatusActive),
+				},
+				parsedArn: arn.ARN{
+					Resource: "container-instance/container_instance_ID",
+				},
+				containerInstanceID: "container_instance_ID",
+			},
+		},
+		{
+			name:    "Found container instance DRAINING and running tasks",
+			wantErr: true,
+			args: args{
+				containerInstance: &ecs.ContainerInstance{
+					Status:            aws.String(ecs.ContainerInstanceStatusDraining),
+					RunningTasksCount: aws.Int64(10),
+				},
+				parsedArn: arn.ARN{
+					Resource: "container-instance/container_instance_ID",
+				},
+				containerInstanceID: "container_instance_ID",
+			},
+		},
+		{
+			name:    "Found container instance DRAINING and no longer running tasks",
+			wantErr: false,
+			args: args{
+				containerInstance: &ecs.ContainerInstance{
+					Status:            aws.String(ecs.ContainerInstanceStatusDraining),
+					RunningTasksCount: aws.Int64(0),
+				},
+				parsedArn: arn.ARN{
+					Resource: "container-instance/container_instance_ID",
+				},
+				containerInstanceID: "container_instance_ID",
+			},
+		},
+		{
+			name:    "Not matching container instance ID",
+			wantErr: true,
+			args: args{
+				containerInstance: &ecs.ContainerInstance{},
+				parsedArn: arn.ARN{
+					Resource: "container-instance/another_container_instance_ID",
+				},
+				containerInstanceID: "container_instance_ID",
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if err := checkDrainingContainerInstance(tt.args.containerInstance, tt.args.parsedArn, tt.args.containerInstanceID); (err != nil) != tt.wantErr {
+				t.Errorf("checkDrainingContainerInstance() error = %v, wantErr %v", err, tt.wantErr)
+			}
+		})
+	}
+}
diff --git a/cmd/update-aws-ecs-service/flags_test.go b/cmd/update-aws-ecs-service/flags_test.go
index 8e4ba09..19a837e 100644
--- a/cmd/update-aws-ecs-service/flags_test.go
+++ b/cmd/update-aws-ecs-service/flags_test.go
@@ -19,8 +19,7 @@ func TestMapMapMapFlag_Set(t *testing.T) {
 	if err := actualStruct.Set("container2=fluentd=option1=value1"); err != nil {
 		t.Fatal(err)
 	}
-	var expectedStruct mapMapMapFlag
-	expectedStruct = map[string]map[string]map[string]string{
+	var expectedStruct mapMapMapFlag = map[string]map[string]map[string]string{
 		"container1": {
 			"awslogs": {
 				"region":   "us-west-2",
diff --git a/ecs.go b/ecs.go
index 54801d7..62429b4 100644
--- a/ecs.go
+++ b/ecs.go
@@ -3,10 +3,12 @@ package awsecs
 import (
 	"errors"
 	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/arn"
 	"github.com/aws/aws-sdk-go/service/ecs"
 	"github.com/cenkalti/backoff"
 	"log"
 	"reflect"
+	"strings"
 )
 
 var (
@@ -22,6 +24,8 @@ var (
 	ErrServiceNotFound = errors.New("the service does not exist")
 	// ErrServiceDeletedAfterUpdate service was updated and then deleted elsewhere
 	ErrServiceDeletedAfterUpdate = backoff.Permanent(errors.New("the service was deleted after the update"))
+	// ErrContainerInstanceNotFound the container instance was removed from the cluster elsewhere
+	ErrContainerInstanceNotFound = backoff.Permanent(errors.New("container instance not found"))
 )
 
 var (
@@ -163,8 +167,8 @@ func copyTaskDef(api ecs.ECS, taskdef string, imageMap map[string]string, envMap
 	if err != nil {
 		return "", err
 	}
-	arn := tdNew.TaskDefinition.TaskDefinitionArn
-	return *arn, nil
+	taskDefinitionArn := tdNew.TaskDefinition.TaskDefinitionArn
+	return *taskDefinitionArn, nil
 }
 
 func alterService(api ecs.ECS, cluster, service string, imageMap map[string]string, envMaps map[string]map[string]string, secretMaps map[string]map[string]string, logopts map[string]map[string]map[string]string, logsecrets map[string]map[string]map[string]string, taskRole string, desiredCount *int64, taskdef string) (ecs.Service, ecs.Service, error) {
@@ -172,19 +176,51 @@ func alterService(api ecs.ECS, cluster, service string, imageMap map[string]stri
 	if err != nil {
 		return ecs.Service{}, ecs.Service{}, err
 	}
-	for _, svc := range output.Services {
+	copyTaskDefinitionAction := func(sourceTaskDefinition string) (string, error) {
+		return copyTaskDef(api, sourceTaskDefinition, imageMap, envMaps, secretMaps, logopts, logsecrets, taskRole)
+	}
+	updateAction := func(newTaskDefinition *string, desiredCount *int64) (*ecs.UpdateServiceOutput, error) {
+		updateServiceInput := &ecs.UpdateServiceInput{
+			Cluster:            aws.String(cluster),
+			Service:            aws.String(service),
+			TaskDefinition:     newTaskDefinition,
+			DesiredCount:       desiredCount,
+			ForceNewDeployment: aws.Bool(true),
+		}
+		return api.UpdateService(updateServiceInput)
+	}
+	return findAndUpdateService(output, cluster, service, taskdef, desiredCount, copyTaskDefinitionAction, updateAction)
+}
+
+func findAndUpdateService(output *ecs.DescribeServicesOutput, cluster, service, taskDefinition string, desiredCount *int64, copyTdAction func(string) (string, error), updateSvcAction func(*string, *int64) (*ecs.UpdateServiceOutput, error)) (ecs.Service, ecs.Service, error) {
+	if len(output.Services) == 0 {
+		return ecs.Service{}, ecs.Service{}, ErrServiceNotFound
+	}
+	svc := output.Services[0]
+	clusterArn := *svc.ClusterArn
+	parsedClusterArn, err := arn.Parse(clusterArn)
+	if err != nil {
+		return ecs.Service{}, ecs.Service{}, err
+	}
+	return updateService(parsedClusterArn, svc, cluster, service, taskDefinition, desiredCount, copyTdAction, updateSvcAction)
+}
+
+func updateService(parsedClusterArn arn.ARN, svc *ecs.Service, cluster, service, td string, desiredCount *int64, copyTdAction func(string) (string, error), updateSvcAction func(*string, *int64) (*ecs.UpdateServiceOutput, error)) (ecs.Service, ecs.Service, error) {
+	clusterNameFound := strings.TrimPrefix(parsedClusterArn.Resource, "cluster/")
+	serviceNameFound := *svc.ServiceName
+	if clusterNameFound == cluster && serviceNameFound == service {
 		srcTaskDef := svc.TaskDefinition
-		if taskdef != "" {
-			srcTaskDef = &taskdef
+		if td != "" {
+			srcTaskDef = &td
 		}
-		newTd, err := copyTaskDef(api, *srcTaskDef, imageMap, envMaps, secretMaps, logopts, logsecrets, taskRole)
+		newTd, err := copyTdAction(*srcTaskDef)
 		if err != nil {
 			return *svc, ecs.Service{}, err
 		}
 		if desiredCount == nil {
 			desiredCount = svc.DesiredCount
 		}
-		updated, err := api.UpdateService(&ecs.UpdateServiceInput{Cluster: aws.String(cluster), Service: aws.String(service), TaskDefinition: aws.String(newTd), DesiredCount: desiredCount, ForceNewDeployment: aws.Bool(true)})
+		updated, err := updateSvcAction(aws.String(newTd), desiredCount)
 		if err != nil {
 			return *svc, ecs.Service{}, err
 		}
diff --git a/ecs_test.go b/ecs_test.go
new file mode 100644
index 0000000..ad16bf6
--- /dev/null
+++ b/ecs_test.go
@@ -0,0 +1,160 @@
+package awsecs
+
+import (
+	"errors"
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/arn"
+	"github.com/aws/aws-sdk-go/service/ecs"
+	"reflect"
+	"testing"
+)
+
+func TestUpdateService(t *testing.T) {
+	type args struct {
+		parsedClusterArn arn.ARN
+		svc              *ecs.Service
+		cluster          string
+		service          string
+		td               string
+		desiredCount     *int64
+		copyTdAction     func(string) (string, error)
+		updateSvcAction  func(*string, *int64) (*ecs.UpdateServiceOutput, error)
+	}
+	tests := []struct {
+		name         string
+		wantErr      bool
+		beforeUpdate ecs.Service
+		afterUpdate  ecs.Service
+		args         args
+	}{
+		{
+			name:    "On copy error I want error",
+			wantErr: true,
+			beforeUpdate: ecs.Service{
+				ServiceName: aws.String("my-service"),
+			},
+			afterUpdate: ecs.Service{},
+			args: args{
+				parsedClusterArn: arn.ARN{
+					Resource: "cluster/my-cluster",
+				},
+				svc: &ecs.Service{
+					ServiceName: aws.String("my-service"),
+				},
+				cluster:      "my-cluster",
+				service:      "my-service",
+				td:           "task:1",
+				desiredCount: aws.Int64(1),
+				copyTdAction: func(string) (string, error) {
+					return "", errors.New("failed to copy")
+				},
+				updateSvcAction: nil,
+			},
+		},
+		{
+			name:    "On update error I want error",
+			wantErr: true,
+			beforeUpdate: ecs.Service{
+				ServiceName: aws.String("my-service"),
+			},
+			afterUpdate: ecs.Service{},
+			args: args{
+				parsedClusterArn: arn.ARN{
+					Resource: "cluster/my-cluster",
+				},
+				svc: &ecs.Service{
+					ServiceName: aws.String("my-service"),
+				},
+				cluster:      "my-cluster",
+				service:      "my-service",
+				td:           "task:1",
+				desiredCount: aws.Int64(1),
+				copyTdAction: func(string) (string, error) {
+					return "task:2", nil
+				},
+				updateSvcAction: func(*string, *int64) (*ecs.UpdateServiceOutput, error) {
+					return nil, errors.New("failed to update")
+				},
+			},
+		},
+		{
+			name:         "On non matching cluster I want error",
+			wantErr:      true,
+			beforeUpdate: ecs.Service{},
+			afterUpdate:  ecs.Service{},
+			args: args{
+				parsedClusterArn: arn.ARN{
+					Resource: "cluster/my-cluster",
+				},
+				svc: &ecs.Service{
+					ServiceName: aws.String("my-service"),
+				},
+				cluster: "my-other-cluster",
+				service: "my-service",
+			},
+		},
+		{
+			name:         "On non matching service I want error",
+			wantErr:      true,
+			beforeUpdate: ecs.Service{},
+			afterUpdate:  ecs.Service{},
+			args: args{
+				parsedClusterArn: arn.ARN{
+					Resource: "cluster/my-cluster",
+				},
+				svc: &ecs.Service{
+					ServiceName: aws.String("my-service"),
+				},
+				cluster: "my-cluster",
+				service: "my-other-service",
+			},
+		},
+		{
+			name:    "Check before and after update",
+			wantErr: false,
+			beforeUpdate: ecs.Service{
+				ServiceName: aws.String("my-service"),
+			},
+			afterUpdate: ecs.Service{
+				TaskDefinition: aws.String("task:2"),
+			},
+			args: args{
+				parsedClusterArn: arn.ARN{
+					Resource: "cluster/my-cluster",
+				},
+				svc: &ecs.Service{
+					ServiceName: aws.String("my-service"),
+				},
+				cluster:      "my-cluster",
+				service:      "my-service",
+				td:           "task:1",
+				desiredCount: nil,
+				copyTdAction: func(s string) (string, error) {
+					return "task:2", nil
+				},
+				updateSvcAction: func(s *string, i *int64) (*ecs.UpdateServiceOutput, error) {
+					return &ecs.UpdateServiceOutput{
+						Service: &ecs.Service{
+							TaskDefinition: aws.String("task:2"),
+						},
+					}, nil
+				},
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, got1, err := updateService(tt.args.parsedClusterArn, tt.args.svc, tt.args.cluster, tt.args.service, tt.args.td, tt.args.desiredCount, tt.args.copyTdAction, tt.args.updateSvcAction)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("updateService() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if !reflect.DeepEqual(got, tt.beforeUpdate) {
+				t.Errorf("updateService() got = %v, want %v", got, tt.beforeUpdate)
+			}
+			if !reflect.DeepEqual(got1, tt.afterUpdate) {
+				t.Errorf("updateService() got1 = %v, want %v", got1, tt.afterUpdate)
+			}
+		})
+	}
+}