add test for draining case

Author: kathap

handlers/upload_droplet/upload_droplet_test.go

@@ -9,6 +9,7 @@ import (
 	"net/http/httptest"
 	"net/url"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"code.cloudfoundry.org/cc-uploader/ccclient/fake_ccclient"
@@ -28,6 +29,7 @@ var _ = Describe("UploadDroplet", func() {
 		var uploader fake_ccclient.FakeUploader
 		var poller fake_ccclient.FakePoller
 		var logger lager.Logger
+		//var draining int32
 
 		BeforeEach(func() {
 			outgoingResponse = httptest.NewRecorder()
@@ -179,6 +181,27 @@ var _ = Describe("UploadDroplet", func() {
 			})
 		})
 
+		Context("when the system is draining", func() {
+			var dropletUploadHandler http.Handler
+			var draining int32
+			var responseWriter  *httptest.ResponseRecorder
+			BeforeEach(func() {
+				atomic.StoreInt32(&draining, 1)
+				responseWriter = httptest.NewRecorder()
+				var err error
+				incomingRequest, err = http.NewRequest("POST", "http://example.com?cc_droplet_upload_uri=http://some-uri", nil)
+				Expect(err).NotTo(HaveOccurred())
+
+				dropletUploadHandler = upload_droplet.New(&uploader, &poller, lager.NewLogger("fake-logger"), &sync.WaitGroup{}, &draining)
+			})
+
+			It("returns 503", func() {
+				dropletUploadHandler.ServeHTTP(responseWriter, incomingRequest)
+				Expect(responseWriter.Code).To(Equal(http.StatusServiceUnavailable))
+				Expect(responseWriter.Body.String()).To(ContainSubstring("Service is draining"))
+			})
+		})
+
 		Context("when the requester (client) goes away", func() {
 			var (
 				cancel             context.CancelFunc