Adapt new tests to not using a 2nd token cache but a 2nd credentials …

…future
flyteorg · Sep 22, 2023 · 4a9735b · 4a9735b
1 parent b9d6917
commit 4a9735b
Showing 1 changed file with 106 additions and 72 deletions.
diff --git a/clients/go/admin/auth_interceptor_test.go b/clients/go/admin/auth_interceptor_test.go
@@ -15,6 +15,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
+	"golang.org/x/oauth2"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -292,75 +293,108 @@ func TestMaterializeCredentials(t *testing.T) {
 	})
 }
 
-// type testRoundTripper struct {
-// 	RoundTripFunc func(req *http.Request) (*http.Response, error)
-// }
-
-// func (t *testRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
-// 	return t.RoundTripFunc(req)
-// }
-
-// func TestSetHTTPClientContext(t *testing.T) {
-// 	ctx := context.Background()
-// 	tokenCache := &cache.TokenCacheInMemoryProvider{}
-
-// 	t.Run("no proxy command and no proxy url", func(t *testing.T) {
-// 		cfg := &Config{}
-
-// 		newCtx, err := setHTTPClientContext(ctx, cfg, tokenCache)
-// 		assert.NoError(t, err)
-
-// 		httpClient, ok := newCtx.Value(oauth2.HTTPClient).(*http.Client)
-// 		assert.True(t, ok)
-
-// 		transport, ok := httpClient.Transport.(*http.Transport)
-// 		assert.True(t, ok)
-// 		assert.Nil(t, transport.Proxy)
-// 	})
-
-// 	t.Run("proxy url", func(t *testing.T) {
-// 		cfg := &Config{
-// 			HTTPProxyURL: config.
-// 				URL{URL: url.URL{
-// 				Scheme: "http",
-// 				Host:   "localhost:8080",
-// 			}},
-// 		}
-// 		newCtx, err := setHTTPClientContext(ctx, cfg, tokenCache)
-// 		assert.NoError(t, err)
-
-// 		httpClient, ok := newCtx.Value(oauth2.HTTPClient).(*http.Client)
-// 		assert.True(t, ok)
-
-// 		transport, ok := httpClient.Transport.(*http.Transport)
-// 		assert.True(t, ok)
-// 		assert.NotNil(t, transport.Proxy)
-// 	})
-
-// 	t.Run("proxy command adds proxy-authorization header", func(t *testing.T) {
-// 		cfg := &Config{
-// 			ProxyCommand: []string{"echo", "test-token-http-client"},
-// 		}
-// 		newCtx, err := setHTTPClientContext(ctx, cfg, tokenCache)
-// 		assert.NoError(t, err)
-
-// 		httpClient, ok := newCtx.Value(oauth2.HTTPClient).(*http.Client)
-// 		assert.True(t, ok)
-
-// 		pat, ok := httpClient.Transport.(*proxyAuthTransport)
-// 		assert.True(t, ok)
-
-// 		testRoundTripper := &testRoundTripper{
-// 			RoundTripFunc: func(req *http.Request) (*http.Response, error) {
-// 				// Check if the ProxyAuthorizationHeader is correctly set
-// 				assert.Equal(t, "Bearer test-token-http-client", req.Header.Get(ProxyAuthorizationHeader))
-// 				return &http.Response{StatusCode: http.StatusOK}, nil
-// 			},
-// 		}
-// 		pat.transport = testRoundTripper
-
-// 		req, _ := http.NewRequest("GET", "http://example.com", nil)
-// 		_, err = httpClient.Do(req)
-// 		assert.NoError(t, err)
-// 	})
-// }
+func TestNewProxyAuthInterceptor(t *testing.T) {
+	cfg := &Config{
+		ProxyCommand: []string{"echo", "test-token"},
+	}
+
+	p := NewPerRPCCredentialsFuture()
+
+	interceptor := NewProxyAuthInterceptor(cfg, p)
+
+	ctx := context.Background()
+	method := "/test.method"
+	req := "request"
+	reply := "reply"
+	cc := new(grpc.ClientConn)
+
+	errorInvoker := func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, opts ...grpc.CallOption) error {
+		return errors.New("test error")
+	}
+
+	// Call should return an error and trigger the interceptor to materialize proxy auth credentials
+	err := interceptor(ctx, method, req, reply, cc, errorInvoker)
+	assert.Error(t, err)
+
+	// Check if proxyCredentialsFuture contains a proxy auth header token
+	creds, err := p.Get().GetRequestMetadata(ctx, "")
+	assert.True(t, p.IsInitialized())
+	assert.NoError(t, err)
+	assert.Equal(t, "Bearer test-token", creds[ProxyAuthorizationHeader])
+}
+
+type testRoundTripper struct {
+	RoundTripFunc func(req *http.Request) (*http.Response, error)
+}
+
+func (t *testRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
+	return t.RoundTripFunc(req)
+}
+
+func TestSetHTTPClientContext(t *testing.T) {
+	ctx := context.Background()
+
+	t.Run("no proxy command and no proxy url", func(t *testing.T) {
+		cfg := &Config{}
+
+		newCtx, err := setHTTPClientContext(ctx, cfg, nil)
+		assert.NoError(t, err)
+
+		httpClient, ok := newCtx.Value(oauth2.HTTPClient).(*http.Client)
+		assert.True(t, ok)
+
+		transport, ok := httpClient.Transport.(*http.Transport)
+		assert.True(t, ok)
+		assert.Nil(t, transport.Proxy)
+	})
+
+	t.Run("proxy url", func(t *testing.T) {
+		cfg := &Config{
+			HTTPProxyURL: config.
+				URL{URL: url.URL{
+				Scheme: "http",
+				Host:   "localhost:8080",
+			}},
+		}
+		newCtx, err := setHTTPClientContext(ctx, cfg, nil)
+		assert.NoError(t, err)
+
+		httpClient, ok := newCtx.Value(oauth2.HTTPClient).(*http.Client)
+		assert.True(t, ok)
+
+		transport, ok := httpClient.Transport.(*http.Transport)
+		assert.True(t, ok)
+		assert.NotNil(t, transport.Proxy)
+	})
+
+	t.Run("proxy command adds proxy-authorization header", func(t *testing.T) {
+		cfg := &Config{
+			ProxyCommand: []string{"echo", "test-token-http-client"},
+		}
+
+		p := NewPerRPCCredentialsFuture()
+		MaterializeProxyAuthCredentials(ctx, cfg, p)
+
+		newCtx, err := setHTTPClientContext(ctx, cfg, p)
+		assert.NoError(t, err)
+
+		httpClient, ok := newCtx.Value(oauth2.HTTPClient).(*http.Client)
+		assert.True(t, ok)
+
+		pat, ok := httpClient.Transport.(*proxyAuthTransport)
+		assert.True(t, ok)
+
+		testRoundTripper := &testRoundTripper{
+			RoundTripFunc: func(req *http.Request) (*http.Response, error) {
+				// Check if the ProxyAuthorizationHeader is correctly set
+				assert.Equal(t, "Bearer test-token-http-client", req.Header.Get(ProxyAuthorizationHeader))
+				return &http.Response{StatusCode: http.StatusOK}, nil
+			},
+		}
+		pat.transport = testRoundTripper
+
+		req, _ := http.NewRequest("GET", "http://example.com", nil)
+		_, err = httpClient.Do(req)
+		assert.NoError(t, err)
+	})
+}