Merge branch '21535-multi-wf-delete'

[arvados.git] / sdk / go / arvadostest / proxy.go

diff --git a/sdk/go/arvadostest/proxy.go b/sdk/go/arvadostest/proxy.go
index 48700d8b186d8fd4d104ed820a6d3060772a86bc..85d433089aa8b7ba746660fafb2aa64457b83067 100644 (file)
--- a/sdk/go/arvadostest/proxy.go
+++ b/sdk/go/arvadostest/proxy.go
@@ -11,6 +11,7 @@ import (
        "net/http/httptest"
        "net/http/httputil"
        "net/url"
+       "sync"
        "time"
 
        "git.arvados.org/arvados.git/sdk/go/arvados"
@@ -26,6 +27,12 @@ type Proxy struct {
 
        // A dump of each request that has been proxied.
        RequestDumps [][]byte
+
+       // If non-nil, func will be called on each incoming request
+       // before proxying it.
+       Director func(*http.Request)
+
+       wg sync.WaitGroup
 }
 
 // NewProxy returns a new Proxy that saves a dump of each reqeust
@@ -62,11 +69,25 @@ func NewProxy(c *check.C, svc arvados.Service) *Proxy {
                Server: srv,
                URL:    u,
        }
+       var mtx sync.Mutex
        rp.Director = func(r *http.Request) {
+               proxy.wg.Add(1)
+               defer proxy.wg.Done()
+               if proxy.Director != nil {
+                       proxy.Director(r)
+               }
                dump, _ := httputil.DumpRequest(r, true)
+               mtx.Lock()
                proxy.RequestDumps = append(proxy.RequestDumps, dump)
+               mtx.Unlock()
                r.URL.Scheme = target.Scheme
                r.URL.Host = target.Host
        }
        return proxy
 }
+
+// Wait waits until all of the proxied requests that have been sent to
+// Director() have also been recorded in RequestDumps.
+func (proxy *Proxy) Wait() {
+       proxy.wg.Wait()
+}