Rewrite TestConcurrentClientRequest to not depend on internal BackendManagers

Author: tallclair

tests/concurrent_client_request_test.go

@@ -19,39 +19,34 @@ package tests
 import (
 	"bytes"
 	"context"
-	"fmt"
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"strconv"
 	"sync"
 	"testing"
+	"time"
 
 	"google.golang.org/grpc"
 	"k8s.io/apimachinery/pkg/util/wait"
 	"sigs.k8s.io/apiserver-network-proxy/konnectivity-client/pkg/client"
-	"sigs.k8s.io/apiserver-network-proxy/pkg/server"
-	"sigs.k8s.io/apiserver-network-proxy/proto/header"
 )
 
 type simpleServer struct {
-	receivedSecondReq chan struct{}
+	mu sync.Mutex
 }
 
-// ServeHTTP blocks the response to the request whose body is "1" until a
-// request whose body is "2" is handled.
 func (s *simpleServer) ServeHTTP(w http.ResponseWriter, req *http.Request) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	time.Sleep(time.Millisecond)
+
 	bytes, err := io.ReadAll(req.Body)
 	if err != nil {
 		w.Write([]byte(err.Error()))
 	}
-	if string(bytes) == "2" {
-		close(s.receivedSecondReq)
-		w.Write([]byte("2"))
-	}
-	if string(bytes) == "1" {
-		<-s.receivedSecondReq
-		w.Write([]byte("1"))
-	}
+	w.Write(bytes)
 }
 
 // TODO: test http-connect as well.
@@ -70,119 +65,45 @@ func getTestClient(front string, t *testing.T) *http.Client {
 	}
 }
 
-// singleTimeManager makes sure that a backend only serves one request.
-type singleTimeManager struct {
-	mu       sync.Mutex
-	backends map[string]server.Backend
-	used     map[string]struct{}
-}
-
-func (s *singleTimeManager) AddBackend(agentID string, _ header.IdentifierType, backend server.Backend) {
-	s.mu.Lock()
-	defer s.mu.Unlock()
-	s.backends[agentID] = backend
-}
-
-func (s *singleTimeManager) RemoveBackend(agentID string, _ header.IdentifierType, backend server.Backend) {
-	s.mu.Lock()
-	defer s.mu.Unlock()
-	v, ok := s.backends[agentID]
-	if !ok {
-		panic(fmt.Errorf("no backends found for %s", agentID))
-	}
-	if v != backend {
-		panic(fmt.Errorf("recorded backend %v does not match %v", &v, backend))
-	}
-	delete(s.backends, agentID)
-}
-
-func (s *singleTimeManager) Backend(_ context.Context) (server.Backend, error) {
-	s.mu.Lock()
-	defer s.mu.Unlock()
-	for k, v := range s.backends {
-		if _, ok := s.used[k]; !ok {
-			s.used[k] = struct{}{}
-			return v, nil
-		}
-	}
-	return nil, fmt.Errorf("cannot find backend to a new agent")
-}
-
-func (s *singleTimeManager) GetBackend(agentID string) server.Backend {
-	return nil
-}
-
-func (s *singleTimeManager) NumBackends() int {
-	return 0
-}
-
-func newSingleTimeGetter(m *server.DefaultBackendManager) *singleTimeManager {
-	return &singleTimeManager{
-		used:     make(map[string]struct{}),
-		backends: make(map[string]server.Backend),
-	}
-}
-
-var _ server.BackendManager = &singleTimeManager{}
-
-func (s *singleTimeManager) Ready() (bool, string) {
-	return true, ""
-}
-
 func TestConcurrentClientRequest(t *testing.T) {
-	t.Skip() // FIXME: figure out how to run this without overriding the BackendManagers
-	s := httptest.NewServer(&simpleServer{receivedSecondReq: make(chan struct{})})
+	const numConcurrentRequests = 100
+	s := httptest.NewServer(&simpleServer{})
 	defer s.Close()
 
 	ps := runGRPCProxyServerWithServerCount(t, 1)
 	defer ps.Stop()
-	// ps.BackendManagers = []server.BackendManager{newSingleTimeGetter(server.NewDefaultBackendManager())} FIXME
 
 	// Run two agents
-	ai1 := runAgent(t, ps.AgentAddr())
-	ai2 := runAgent(t, ps.AgentAddr())
-	defer ai1.Stop()
-	defer ai2.Stop()
-	waitForConnectedServerCount(t, 1, ai1)
-	waitForConnectedServerCount(t, 1, ai2)
-
-	client1 := getTestClient(ps.FrontAddr(), t)
-	client2 := getTestClient(ps.FrontAddr(), t)
+	a1 := runAgent(t, ps.AgentAddr())
+	a2 := runAgent(t, ps.AgentAddr())
+	defer a1.Stop()
+	defer a2.Stop()
+	waitForConnectedServerCount(t, 1, a1)
+	waitForConnectedServerCount(t, 1, a2)
+
 	var wg sync.WaitGroup
-	wg.Add(2)
-	go func() {
-		defer wg.Done()
-		r, err := client1.Post(s.URL, "text/plain", bytes.NewBufferString("1"))
-		if err != nil {
-			t.Error(err)
-			return
-		}
-		data, err := io.ReadAll(r.Body)
-		if err != nil {
-			t.Error(err)
-		}
-		r.Body.Close()
-
-		if string(data) != "1" {
-			t.Errorf("expect %v; got %v", "1", string(data))
-		}
-	}()
-	go func() {
-		defer wg.Done()
-		r, err := client2.Post(s.URL, "text/plain", bytes.NewBufferString("2"))
-		if err != nil {
-			t.Error(err)
-			return
-		}
-		data, err := io.ReadAll(r.Body)
-		if err != nil {
-			t.Error(err)
-		}
-		r.Body.Close()
-
-		if string(data) != "2" {
-			t.Errorf("expect %v; got %v", "2", string(data))
-		}
-	}()
+	wg.Add(numConcurrentRequests)
+	for i := 0; i < numConcurrentRequests; i++ {
+		id := i
+		go func() {
+			defer wg.Done()
+			client1 := getTestClient(ps.FrontAddr(), t)
+
+			r, err := client1.Post(s.URL, "text/plain", bytes.NewBufferString(strconv.Itoa(id)))
+			if err != nil {
+				t.Error(err)
+				return
+			}
+			data, err := io.ReadAll(r.Body)
+			if err != nil {
+				t.Error(err)
+			}
+			r.Body.Close()
+
+			if string(data) != strconv.Itoa(id) {
+				t.Errorf("expect %d; got %s", id, string(data))
+			}
+		}()
+	}
 	wg.Wait()
 }