mockgcp: serve multiple hosts, including variables

gemmahou · Aug 8, 2024 · 7a2507f · 7a2507f
1 parent 395b43f
commit 7a2507f
Show file tree

Hide file tree

Showing 39 changed files with 128 additions and 90 deletions.
diff --git a/mockgcp/mock_http_roundtrip.go b/mockgcp/mock_http_roundtrip.go
@@ -23,6 +23,7 @@ import (
 	"log"
 	"net"
 	"net/http"
+	"regexp"
 	"strings"
 	"testing"
 
@@ -78,11 +79,24 @@ type mockRoundTripper struct {
 	grpcConnection *grpc.ClientConn
 	grpcListener   net.Listener
 
-	hosts map[string]http.Handler
-
 	iamPolicies *mockIAMPolicies
 
-	services []MockService
+	services []registeredService
+}
+
+type registeredService struct {
+	hostRegexes []*regexp.Regexp
+	handler     http.Handler
+	impl        MockService
+}
+
+func (h *registeredService) MatchesHost(host string) (http.Handler, bool) {
+	for _, hostRegex := range h.hostRegexes {
+		if hostRegex.MatchString(host) {
+			return h.handler, true
+		}
+	}
+	return nil, false
 }
 
 // MockService is the interface implemented by all services
@@ -93,8 +107,9 @@ type MockService interface {
 	// NewHTTPMux creates an HTTP mux for serving http traffic
 	NewHTTPMux(ctx context.Context, conn *grpc.ClientConn) (http.Handler, error)
 
-	// ExpectedHost is the hostname we serve on e.g. foo.googleapis.com
-	ExpectedHost() string
+	// ExpectedHosts is the hostname(s) we serve on e.g. foo.googleapis.com
+	// We also support patterns like `{region}-foo.googleapis.com`
+	ExpectedHosts() []string
 }
 
 type Interface interface {
@@ -137,8 +152,6 @@ func NewMockRoundTripper(t *testing.T, k8sClient client.Client, storage storage.
 	var serverOpts []grpc.ServerOption
 	server := grpc.NewServer(serverOpts...)
 
-	mockRoundTripper.hosts = make(map[string]http.Handler)
-
 	var services []MockService
 
 	services = append(services, resourcemanagerService)
@@ -179,7 +192,6 @@ func NewMockRoundTripper(t *testing.T, k8sClient client.Client, storage storage.
 	services = append(services, mockcontaineranalysis.New(env, storage))
 	services = append(services, mockdataform.New(env, storage))
 	services = append(services, mockbigqueryconnection.New(env, storage))
-	mockRoundTripper.services = services
 
 	for _, service := range services {
 		service.Register(server)
@@ -216,7 +228,15 @@ func NewMockRoundTripper(t *testing.T, k8sClient client.Client, storage storage.
 		if err != nil {
 			t.Fatalf("error building mux: %v", err)
 		}
-		mockRoundTripper.hosts[service.ExpectedHost()] = mux
+		var hostRegexes []*regexp.Regexp
+		for _, host := range service.ExpectedHosts() {
+			hostRegexes = append(hostRegexes, toHostRegex(host))
+		}
+		mockRoundTripper.services = append(mockRoundTripper.services, registeredService{
+			impl:        service,
+			hostRegexes: hostRegexes,
+			handler:     mux,
+		})
 	}
 
 	mockRoundTripper.iamPolicies = newMockIAMPolicies()
@@ -226,11 +246,11 @@ func NewMockRoundTripper(t *testing.T, k8sClient client.Client, storage storage.
 
 func (m *mockRoundTripper) RunTestCommand(ctx context.Context, serviceName string, command string) error {
 	for _, service := range m.services {
-		if service.ExpectedHost() != serviceName {
+		if _, match := service.MatchesHost(serviceName); !match {
 			continue
 		}
 
-		supportsTestCommands, ok := service.(SupportsTestCommands)
+		supportsTestCommands, ok := service.impl.(SupportsTestCommands)
 		if !ok {
 			return fmt.Errorf("service %T does not support test commands", service)
 		}
@@ -251,6 +271,19 @@ func (m *mockRoundTripper) NewGRPCConnection(ctx context.Context) *grpc.ClientCo
 	return conn
 }
 
+func toHostRegex(host string) *regexp.Regexp {
+	r := regexp.MustCompile(`{[^}]+}`)
+
+	tokens := strings.Split(host, ".")
+	for i, token := range tokens {
+		token = r.ReplaceAllStringFunc(token, func(match string) string {
+			return "[^.]*"
+		})
+		tokens[i] = token
+	}
+	return regexp.MustCompile("^" + strings.Join(tokens, `\.`) + "$")
+}
+
 func (m *mockRoundTripper) prefilterRequest(req *http.Request) error {
 	if req.Body != nil {
 		var requestBody bytes.Buffer
@@ -389,7 +422,14 @@ func (m *mockRoundTripper) RoundTrip(req *http.Request) (*http.Response, error)
 		return m.roundTripIAMPolicy(req)
 	}
 
-	mux := m.hosts[req.Host]
+	var mux http.Handler
+	for _, service := range m.services {
+		m, found := service.MatchesHost(req.Host)
+		if found {
+			mux = m
+			break
+		}
+	}
 	if mux != nil {
 		if err := m.prefilterRequest(req); err != nil {
 			return nil, err

diff --git a/mockgcp/mockaiplatform/service.go b/mockgcp/mockaiplatform/service.go
@@ -46,9 +46,8 @@ func New(env *common.MockEnvironment, storage storage.Storage) *MockService {
 	return s
 }
 
-func (s *MockService) ExpectedHost() string {
-	// TODO: Support more endpoints
-	return "us-central1-aiplatform.googleapis.com"
+func (s *MockService) ExpectedHosts() []string {
+	return []string{"{region}-aiplatform.googleapis.com"}
 }
 
 func (s *MockService) Register(grpcServer *grpc.Server) {

diff --git a/mockgcp/mockalloydb/service.go b/mockgcp/mockalloydb/service.go
@@ -46,8 +46,8 @@ func New(env *common.MockEnvironment, storage storage.Storage) *MockService {
 	return s
 }
 
-func (s *MockService) ExpectedHost() string {
-	return "alloydb.googleapis.com"
+func (s *MockService) ExpectedHosts() []string {
+	return []string{"alloydb.googleapis.com"}
 }
 
 func (s *MockService) Register(grpcServer *grpc.Server) {

diff --git a/mockgcp/mockapikeys/service.go b/mockgcp/mockapikeys/service.go
@@ -48,8 +48,8 @@ func New(env *common.MockEnvironment, storage storage.Storage) *MockService {
 	return s
 }
 
-func (s *MockService) ExpectedHost() string {
-	return "apikeys.googleapis.com"
+func (s *MockService) ExpectedHosts() []string {
+	return []string{"apikeys.googleapis.com"}
 }
 
 func (s *MockService) Register(grpcServer *grpc.Server) {

diff --git a/mockgcp/mockartifactregistry/service.go b/mockgcp/mockartifactregistry/service.go
@@ -46,8 +46,8 @@ func New(env *common.MockEnvironment, storage storage.Storage) *MockService {
 	return s
 }
 
-func (s *MockService) ExpectedHost() string {
-	return "artifactregistry.googleapis.com"
+func (s *MockService) ExpectedHosts() []string {
+	return []string{"artifactregistry.googleapis.com"}
 }
 
 func (s *MockService) Register(grpcServer *grpc.Server) {

diff --git a/mockgcp/mockbigquery/service.go b/mockgcp/mockbigquery/service.go
@@ -46,8 +46,8 @@ func New(env *common.MockEnvironment, storage storage.Storage) *MockService {
 	return s
 }
 
-func (s *MockService) ExpectedHost() string {
-	return "bigquery.googleapis.com"
+func (s *MockService) ExpectedHosts() []string {
+	return []string{"bigquery.googleapis.com"}
 }
 
 func (s *MockService) Register(grpcServer *grpc.Server) {

diff --git a/mockgcp/mockbigqueryconnection/service.go b/mockgcp/mockbigqueryconnection/service.go
@@ -45,8 +45,8 @@ func New(env *common.MockEnvironment, storage storage.Storage) *MockService {
 	return s
 }
 
-func (s *MockService) ExpectedHost() string {
-	return "bigqueryconnection.googleapis.com"
+func (s *MockService) ExpectedHosts() []string {
+	return []string{"bigqueryconnection.googleapis.com"}
 }
 
 func (s *MockService) Register(grpcServer *grpc.Server) {

diff --git a/mockgcp/mockbigtable/service.go b/mockgcp/mockbigtable/service.go
@@ -50,8 +50,8 @@ func New(env *common.MockEnvironment, storage storage.Storage) *MockService {
 	return s
 }
 
-func (s *MockService) ExpectedHost() string {
-	return "bigtableadmin.googleapis.com"
+func (s *MockService) ExpectedHosts() []string {
+	return []string{"bigtableadmin.googleapis.com"}
 }
 
 func (s *MockService) Register(grpcServer *grpc.Server) {

diff --git a/mockgcp/mockbilling/service.go b/mockgcp/mockbilling/service.go
@@ -48,8 +48,8 @@ func New(env *common.MockEnvironment, storage storage.Storage) *MockService {
 	return s
 }
 
-func (s *MockService) ExpectedHost() string {
-	return "cloudbilling.googleapis.com"
+func (s *MockService) ExpectedHosts() []string {
+	return []string{"cloudbilling.googleapis.com"}
 }
 
 func (s *MockService) Register(grpcServer *grpc.Server) {

diff --git a/mockgcp/mockcertificatemanager/service.go b/mockgcp/mockcertificatemanager/service.go
@@ -48,8 +48,8 @@ func New(env *common.MockEnvironment, storage storage.Storage) *MockService {
 	return s
 }
 
-func (s *MockService) ExpectedHost() string {
-	return "certificatemanager.googleapis.com"
+func (s *MockService) ExpectedHosts() []string {
+	return []string{"certificatemanager.googleapis.com"}
 }
 
 func (s *MockService) Register(grpcServer *grpc.Server) {

diff --git a/mockgcp/mockcloudbuild/service.go b/mockgcp/mockcloudbuild/service.go
@@ -47,8 +47,8 @@ func New(env *common.MockEnvironment, storage storage.Storage) *MockService {
 	return s
 }
 
-func (s *MockService) ExpectedHost() string {
-	return "cloudbuild.googleapis.com"
+func (s *MockService) ExpectedHosts() []string {
+	return []string{"cloudbuild.googleapis.com"}
 }
 
 func (s *MockService) Register(grpcServer *grpc.Server) {

diff --git a/mockgcp/mockcloudfunctions/service.go b/mockgcp/mockcloudfunctions/service.go
@@ -48,8 +48,8 @@ func New(env *common.MockEnvironment, storage storage.Storage) *MockService {
 	return s
 }
 
-func (s *MockService) ExpectedHost() string {
-	return "cloudfunctions.googleapis.com"
+func (s *MockService) ExpectedHosts() []string {
+	return []string{"cloudfunctions.googleapis.com"}
 }
 
 func (s *MockService) Register(grpcServer *grpc.Server) {

diff --git a/mockgcp/mockcloudids/service.go b/mockgcp/mockcloudids/service.go
@@ -47,8 +47,8 @@ func New(env *common.MockEnvironment, storage storage.Storage) *MockService {
 	return s
 }
 
-func (s *MockService) ExpectedHost() string {
-	return "ids.googleapis.com"
+func (s *MockService) ExpectedHosts() []string {
+	return []string{"ids.googleapis.com"}
 }
 
 func (s *MockService) Register(grpcServer *grpc.Server) {

diff --git a/mockgcp/mockcompute/service.go b/mockgcp/mockcompute/service.go
@@ -45,8 +45,8 @@ func New(env *common.MockEnvironment, storage storage.Storage) *MockService {
 	return s
 }
 
-func (s *MockService) ExpectedHost() string {
-	return "compute.googleapis.com"
+func (s *MockService) ExpectedHosts() []string {
+	return []string{"compute.googleapis.com"}
 }
 
 func (s *MockService) Register(grpcServer *grpc.Server) {

diff --git a/mockgcp/mockcontainer/service.go b/mockgcp/mockcontainer/service.go
@@ -41,8 +41,8 @@ func New(env *common.MockEnvironment, storage storage.Storage) *MockService {
 	return s
 }
 
-func (s *MockService) ExpectedHost() string {
-	return "container.googleapis.com"
+func (s *MockService) ExpectedHosts() []string {
+	return []string{"container.googleapis.com"}
 }
 
 func (s *MockService) Register(grpcServer *grpc.Server) {

diff --git a/mockgcp/mockcontaineranalysis/service.go b/mockgcp/mockcontaineranalysis/service.go
@@ -46,8 +46,8 @@ func New(env *common.MockEnvironment, storage storage.Storage) *MockService {
 	return s
 }
 
-func (s *MockService) ExpectedHost() string {
-	return "containeranalysis.googleapis.com"
+func (s *MockService) ExpectedHosts() []string {
+	return []string{"containeranalysis.googleapis.com"}
 }
 
 func (s *MockService) Register(grpcServer *grpc.Server) {

diff --git a/mockgcp/mockdataform/service.go b/mockgcp/mockdataform/service.go
@@ -48,8 +48,8 @@ func New(env *common.MockEnvironment, storage storage.Storage) *MockService {
 	return s
 }
 
-func (s *MockService) ExpectedHost() string {
-	return "dataform.googleapis.com"
+func (s *MockService) ExpectedHosts() []string {
+	return []string{"dataform.googleapis.com"}
 }
 
 func (s *MockService) Register(grpcServer *grpc.Server) {

diff --git a/mockgcp/mockedgecontainer/service.go b/mockgcp/mockedgecontainer/service.go
@@ -53,8 +53,8 @@ func New(env *common.MockEnvironment, storage storage.Storage) *MockService {
 	return s
 }
 
-func (s *MockService) ExpectedHost() string {
-	return "edgecontainer.googleapis.com"
+func (s *MockService) ExpectedHosts() []string {
+	return []string{"edgecontainer.googleapis.com"}
 }
 
 func (s *MockService) Register(grpcServer *grpc.Server) {

diff --git a/mockgcp/mockedgenetwork/service.go b/mockgcp/mockedgenetwork/service.go
@@ -53,8 +53,8 @@ func New(env *common.MockEnvironment, storage storage.Storage) *MockService {
 	return s
 }
 
-func (s *MockService) ExpectedHost() string {
-	return "edgenetwork.googleapis.com"
+func (s *MockService) ExpectedHosts() []string {
+	return []string{"edgenetwork.googleapis.com"}
 }
 
 func (s *MockService) Register(grpcServer *grpc.Server) {

diff --git a/mockgcp/mockgkehub/service.go b/mockgcp/mockgkehub/service.go
@@ -51,8 +51,8 @@ func New(env *common.MockEnvironment, storage storage.Storage) *MockService {
 	return s
 }
 
-func (s *MockService) ExpectedHost() string {
-	return "gkehub.googleapis.com"
+func (s *MockService) ExpectedHosts() []string {
+	return []string{"gkehub.googleapis.com"}
 }
 
 func (s *MockService) Register(grpcServer *grpc.Server) {

diff --git a/mockgcp/mockgkemulticloud/service.go b/mockgcp/mockgkemulticloud/service.go
@@ -48,8 +48,8 @@ func New(env *common.MockEnvironment, storage storage.Storage) *MockService {
 	return s
 }
 
-func (s *MockService) ExpectedHost() string {
-	return "us-west1-gkemulticloud.googleapis.com"
+func (s *MockService) ExpectedHosts() []string {
+	return []string{"{region}-gkemulticloud.googleapis.com"}
 }
 
 func (s *MockService) Register(grpcServer *grpc.Server) {

diff --git a/mockgcp/mockiam/service.go b/mockgcp/mockiam/service.go
@@ -49,8 +49,8 @@ func New(env *common.MockEnvironment, storage storage.Storage) *MockService {
 	return s
 }
 
-func (s *MockService) ExpectedHost() string {
-	return "iam.googleapis.com"
+func (s *MockService) ExpectedHosts() []string {
+	return []string{"iam.googleapis.com"}
 }
 
 func (s *MockService) Register(grpcServer *grpc.Server) {

diff --git a/mockgcp/mockkms/service.go b/mockgcp/mockkms/service.go
@@ -44,8 +44,8 @@ func New(env *common.MockEnvironment, storage storage.Storage) *MockService {
 	return s
 }
 
-func (s *MockService) ExpectedHost() string {
-	return "cloudkms.googleapis.com"
+func (s *MockService) ExpectedHosts() []string {
+	return []string{"cloudkms.googleapis.com"}
 }
 
 func (s *MockService) Register(grpcServer *grpc.Server) {

diff --git a/mockgcp/mocklogging/service.go b/mockgcp/mocklogging/service.go
@@ -44,8 +44,8 @@ func New(env *common.MockEnvironment, storage storage.Storage) *MockService {
 	return s
 }
 
-func (s *MockService) ExpectedHost() string {
-	return "logging.googleapis.com"
+func (s *MockService) ExpectedHosts() []string {
+	return []string{"logging.googleapis.com"}
 }
 
 func (s *MockService) Register(grpcServer *grpc.Server) {

diff --git a/mockgcp/mockmonitoring/service.go b/mockgcp/mockmonitoring/service.go
@@ -45,8 +45,8 @@ func New(env *common.MockEnvironment, storage storage.Storage) *MockService {
 	return s
 }
 
-func (s *MockService) ExpectedHost() string {
-	return "monitoring.googleapis.com"
+func (s *MockService) ExpectedHosts() []string {
+	return []string{"monitoring.googleapis.com"}
 }
 
 func (s *MockService) Register(grpcServer *grpc.Server) {