Use wait instead of TaskExit.

Signed-off-by: Lantao Liu <lantaol@google.com>

pkg/server/container_remove_test.go

@@ -40,6 +40,13 @@ func TestSetContainerRemoving(t *testing.T) {
 			},
 			expectErr: true,
 		},
+		"should return error when container is in starting state": {
+			status: containerstore.Status{
+				CreatedAt: time.Now().UnixNano(),
+				Starting:  true,
+			},
+			expectErr: true,
+		},
 		"should return error when container is in removing state": {
 			status: containerstore.Status{
 				CreatedAt:  time.Now().UnixNano(),
@@ -71,6 +78,8 @@ func TestSetContainerRemoving(t *testing.T) {
 		} else {
 			assert.NoError(t, err)
 			assert.True(t, container.Status.Get().Removing, "removing should be set")
+			assert.NoError(t, resetContainerRemoving(container))
+			assert.False(t, container.Status.Get().Removing, "removing should be reset")
 		}
 	}
 }