Async test for Broker's proxy handler

broker/broker.go

@@ -29,6 +29,7 @@ type BrokerContext struct {
 	// Map keeping track of snowflakeIDs required to match SDP answers from
 	// the second http POST.
 	snowflakeMap map[string]*Snowflake
+	createChan   chan *ProxyRequest
 }
 
 func NewBrokerContext() *BrokerContext {
@@ -37,6 +38,7 @@ func NewBrokerContext() *BrokerContext {
 	return &BrokerContext{
 		snowflakes:   snowflakes,
 		snowflakeMap: make(map[string]*Snowflake),
+		createChan:   make(chan *ProxyRequest),
 	}
 }
 
@@ -54,8 +56,6 @@ type ProxyRequest struct {
 	offerChan chan []byte
 }
 
-var createChan = make(chan *ProxyRequest)
-
 // Create and add a Snowflake to the heap.
 func (sc *BrokerContext) AddSnowflake(id string) *Snowflake {
 	snowflake := new(Snowflake)
@@ -69,9 +69,11 @@ func (sc *BrokerContext) AddSnowflake(id string) *Snowflake {
 }
 
 // Match proxies to clients.
-func (sc *BrokerContext) Broker(proxies <-chan *ProxyRequest) {
+// func (ctx *BrokerContext) Broker(proxies <-chan *ProxyRequest) {
-	for p := range proxies {
+func (ctx *BrokerContext) Broker() {
-		snowflake := sc.AddSnowflake(p.id)
+	// for p := range proxies {
+	for p := range ctx.createChan {
+		snowflake := ctx.AddSnowflake(p.id)
 		// Wait for a client to avail an offer to the snowflake, or timeout
 		// and ask the snowflake to poll later.
 		go func(p *ProxyRequest) {
@@ -81,8 +83,8 @@ func (sc *BrokerContext) Broker(proxies <-chan *ProxyRequest) {
 				p.offerChan <- offer
 			case <-time.After(time.Second * ProxyTimeout):
 				// This snowflake is no longer available to serve clients.
-				heap.Remove(sc.snowflakes, snowflake.index)
+				heap.Remove(ctx.snowflakes, snowflake.index)
-				delete(sc.snowflakeMap, snowflake.id)
+				delete(ctx.snowflakeMap, snowflake.id)
 				p.offerChan <- nil
 			}
 		}(p)
@@ -176,7 +178,7 @@ func proxyHandler(ctx *BrokerContext, w http.ResponseWriter, r *http.Request) {
 	p := new(ProxyRequest)
 	p.id = id
 	p.offerChan = make(chan []byte)
-	createChan <- p
+	ctx.createChan <- p
 
 	// Wait for a client to avail an offer to the snowflake, or timeout
 	// and ask the snowflake to poll later.
@@ -225,7 +227,7 @@ func debugHandler(ctx *BrokerContext, w http.ResponseWriter, r *http.Request) {
 func init() {
 	ctx := NewBrokerContext()
 
-	go ctx.Broker(createChan)
+	go ctx.Broker()
 
 	http.HandleFunc("/robots.txt", robotsTxtHandler)
 	http.HandleFunc("/ip", ipHandler)

broker/snowflake-broker_test.go

@@ -15,7 +15,6 @@ func TestBroker(t *testing.T) {
 		ctx := NewBrokerContext()
 
 		Convey("Adds Snowflake", func() {
-			ctx := NewBrokerContext()
 			So(ctx.snowflakes.Len(), ShouldEqual, 0)
 			So(len(ctx.snowflakeMap), ShouldEqual, 0)
 			ctx.AddSnowflake("foo")
@@ -24,7 +23,6 @@ func TestBroker(t *testing.T) {
 		})
 
 		Convey("Responds to client offers...", func() {
-
 			w := httptest.NewRecorder()
 			data := bytes.NewReader([]byte("test"))
 			r, err := http.NewRequest("POST", "broker.com/client", data)
@@ -55,6 +53,9 @@ func TestBroker(t *testing.T) {
 			})
 
 			Convey("Times out when no proxy responds.", func() {
+				if testing.Short() {
+					return
+				}
 				done := make(chan bool)
 				snowflake := ctx.AddSnowflake("fake")
 				go func() {
@@ -66,7 +67,43 @@ func TestBroker(t *testing.T) {
 				<-done
 				So(w.Code, ShouldEqual, http.StatusGatewayTimeout)
 			})
+		})
 
+		Convey("Responds to proxy polls...", func() {
+			done := make(chan bool)
+			w := httptest.NewRecorder()
+			data := bytes.NewReader([]byte("test"))
+			r, err := http.NewRequest("POST", "broker.com/proxy", data)
+			r.Header.Set("X-Session-ID", "test")
+			So(err, ShouldBeNil)
+
+			Convey("with a client offer if available.", func() {
+				go func(ctx *BrokerContext) {
+					proxyHandler(ctx, w, r)
+					done <- true
+				}(ctx)
+				// Pass a fake client offer to this proxy
+				p := <-ctx.createChan
+				So(p.id, ShouldEqual, "test")
+				p.offerChan <- []byte("fake offer")
+				<-done
+				So(w.Code, ShouldEqual, http.StatusOK)
+				So(w.Body.String(), ShouldEqual, "fake offer")
+			})
+
+			Convey("times out when no client offer is available.", func() {
+				go func(ctx *BrokerContext) {
+					proxyHandler(ctx, w, r)
+					done <- true
+				}(ctx)
+				p := <-ctx.createChan
+				So(p.id, ShouldEqual, "test")
+				// nil means timeout
+				p.offerChan <- nil
+				<-done
+				So(w.Body.String(), ShouldEqual, "")
+				So(w.Code, ShouldEqual, http.StatusGatewayTimeout)
+			})
 		})
 	})
 }