Repository: envoyproxy/ai-gateway
Pull request #398 (2025-02-21 03:14:17)

test: add TestBackendSecurityPolicyController
**Commit Message**

Adding test TestBackendSecurityPolicyController and exporting BackendSecurityPolicy + NewBackendSecurityPolicy


Patch:
diff --git a/internal/controller/backend_security_policy.go b/internal/controller/backend_security_policy.go
index d0c7f47..2f07ee3 100644
--- a/internal/controller/backend_security_policy.go
+++ b/internal/controller/backend_security_policy.go
@@ -29,10 +29,10 @@ import (
 // Temporarily a fixed duration.
 const preRotationWindow = 5 * time.Minute
 
-// backendSecurityPolicyController implements [reconcile.TypedReconciler] for [aigv1a1.BackendSecurityPolicy].
+// BackendSecurityPolicyController implements [reconcile.TypedReconciler] for [aigv1a1.BackendSecurityPolicy].
 //
 // Exported for testing purposes.
-type backendSecurityPolicyController struct {
+type BackendSecurityPolicyController struct {
 	client               client.Client
 	kube                 kubernetes.Interface
 	logger               logr.Logger
@@ -41,8 +41,8 @@ type backendSecurityPolicyController struct {
 	syncAIServiceBackend syncAIServiceBackendFn
 }
 
-func newBackendSecurityPolicyController(client client.Client, kube kubernetes.Interface, logger logr.Logger, syncAIServiceBackend syncAIServiceBackendFn) *backendSecurityPolicyController {
-	return &backendSecurityPolicyController{
+func NewBackendSecurityPolicyController(client client.Client, kube kubernetes.Interface, logger logr.Logger, syncAIServiceBackend syncAIServiceBackendFn) *BackendSecurityPolicyController {
+	return &BackendSecurityPolicyController{
 		client:               client,
 		kube:                 kube,
 		logger:               logger,
@@ -52,7 +52,7 @@ func newBackendSecurityPolicyController(client client.Client, kube kubernetes.In
 }
 
 // Reconcile implements the [reconcile.TypedReconciler] for [aigv1a1.BackendSecurityPolicy].
-func (c *backendSecurityPolicyController) Reconcile(ctx context.Context, req ctrl.Request) (res ctrl.Result, err error) {
+func (c *BackendSecurityPolicyController) Reconcile(ctx context.Context, req ctrl.Request) (res ctrl.Result, err error) {
 	var backendSecurityPolicy aigv1a1.BackendSecurityPolicy
 	if err = c.client.Get(ctx, req.NamespacedName, &backendSecurityPolicy); err != nil {
 		if apierrors.IsNotFound(err) {
@@ -103,7 +103,7 @@ func (c *backendSecurityPolicyController) Reconcile(ctx context.Context, req ctr
 }
 
 // rotateCredential rotates the credentials using the access token from OIDC provider and return the requeue time for next rotation.
-func (c *backendSecurityPolicyController) rotateCredential(ctx context.Context, policy *aigv1a1.BackendSecurityPolicy, oidcCreds egv1a1.OIDC, rotator rotators.Rotator) (time.Duration, error) {
+func (c *BackendSecurityPolicyController) rotateCredential(ctx context.Context, policy *aigv1a1.BackendSecurityPolicy, oidcCreds egv1a1.OIDC, rotator rotators.Rotator) (time.Duration, error) {
 	bspKey := backendSecurityPolicyKey(policy.Namespace, policy.Name)
 
 	var err error
@@ -152,7 +152,7 @@ func backendSecurityPolicyKey(namespace, name string) string {
 	return fmt.Sprintf("%s.%s", name, namespace)
 }
 
-func (c *backendSecurityPolicyController) syncBackendSecurityPolicy(ctx context.Context, bsp *aigv1a1.BackendSecurityPolicy) error {
+func (c *BackendSecurityPolicyController) syncBackendSecurityPolicy(ctx context.Context, bsp *aigv1a1.BackendSecurityPolicy) error {
 	key := backendSecurityPolicyKey(bsp.Namespace, bsp.Name)
 	var aiServiceBackends aigv1a1.AIServiceBackendList
 	err := c.client.List(ctx, &aiServiceBackends, client.MatchingFields{k8sClientIndexBackendSecurityPolicyToReferencingAIServiceBackend: key})
diff --git a/internal/controller/controller.go b/internal/controller/controller.go
index e73a579..722ebd0 100644
--- a/internal/controller/controller.go
+++ b/internal/controller/controller.go
@@ -104,7 +104,7 @@ func StartControllers(ctx context.Context, config *rest.Config, logger logr.Logg
 		return fmt.Errorf("failed to create controller for AIServiceBackend: %w", err)
 	}
 
-	backendSecurityPolicyC := newBackendSecurityPolicyController(c, kubernetes.NewForConfigOrDie(config), logger.
+	backendSecurityPolicyC := NewBackendSecurityPolicyController(c, kubernetes.NewForConfigOrDie(config), logger.
 		WithName("backend-security-policy"), backendC.syncAIServiceBackend)
 	if err = ctrl.NewControllerManagedBy(mgr).
 		For(&aigv1a1.BackendSecurityPolicy{}).
diff --git a/internal/extproc/chatcompletion_processor.go b/internal/extproc/chatcompletion_processor.go
index 56f0911..6cbfaf0 100644
--- a/internal/extproc/chatcompletion_processor.go
+++ b/internal/extproc/chatcompletion_processor.go
@@ -10,16 +10,19 @@ import (
 	"compress/gzip"
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"log/slog"
 
 	corev3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
 	extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
+	typev3 "github.com/envoyproxy/go-control-plane/envoy/type/v3"
 	"google.golang.org/protobuf/types/known/structpb"
 
 	"github.com/envoyproxy/ai-gateway/filterapi"
 	"github.com/envoyproxy/ai-gateway/internal/apischema/openai"
+	"github.com/envoyproxy/ai-gateway/internal/extproc/router"
 	"github.com/envoyproxy/ai-gateway/internal/extproc/translator"
 	"github.com/envoyproxy/ai-gateway/internal/llmcostcel"
 )
@@ -84,6 +87,17 @@ func (c *chatCompletionProcessor) ProcessRequestBody(ctx context.Context, rawBod
 	c.requestHeaders[c.config.modelNameHeaderKey] = model
 	b, err := c.config.router.Calculate(c.requestHeaders)
 	if err != nil {
+		if errors.Is(err, router.ErrNoMatchingRule) {
+			return &extprocv3.ProcessingResponse{
+				Response: &extprocv3.ProcessingResponse_ImmediateResponse{
+					ImmediateResponse: &extprocv3.ImmediateResponse{
+						Status: &typev3.HttpStatus{Code: typev3.StatusCode_NotFound},
+						Body:   []byte(err.Error()),
+					},
+				},
+			}, nil
+		}
+
 		return nil, fmt.Errorf("failed to calculate route: %w", err)
 	}
 	c.logger.Info("Selected backend", "backend", b.Name)
diff --git a/internal/extproc/router/router.go b/internal/extproc/router/router.go
index d0507cf..759663e 100644
--- a/internal/extproc/router/router.go
+++ b/internal/extproc/router/router.go
@@ -15,6 +15,9 @@ import (
 	"github.com/envoyproxy/ai-gateway/filterapi/x"
 )
 
+// ErrNoMatchingRule is the error when no matching rule is found.
+var ErrNoMatchingRule = errors.New("no matching rule found")
+
 // router implements [x.Router].
 type router struct {
 	rules []filterapi.RouteRule
@@ -46,7 +49,7 @@ func (r *router) Calculate(headers map[string]string) (backend *filterapi.Backen
 		}
 	}
 	if rule == nil {
-		return nil, errors.New("no matching rule found")
+		return nil, ErrNoMatchingRule
 	}
 	return r.selectBackendFromRule(rule), nil
 }


Test patch:
diff --git a/internal/controller/backend_security_policy_test.go b/internal/controller/backend_security_policy_test.go
index 521f520..7d71965 100644
--- a/internal/controller/backend_security_policy_test.go
+++ b/internal/controller/backend_security_policy_test.go
@@ -40,7 +40,7 @@ import (
 func TestBackendSecurityController_Reconcile(t *testing.T) {
 	syncFn := internaltesting.NewSyncFnImpl[aigv1a1.AIServiceBackend]()
 	fakeClient := requireNewFakeClientWithIndexes(t)
-	c := newBackendSecurityPolicyController(fakeClient, fake2.NewClientset(), ctrl.Log, syncFn.Sync)
+	c := NewBackendSecurityPolicyController(fakeClient, fake2.NewClientset(), ctrl.Log, syncFn.Sync)
 	backendSecurityPolicyName := "mybackendSecurityPolicy"
 	namespace := "default"
 
@@ -103,7 +103,7 @@ func (m *mockSTSClient) AssumeRoleWithWebIdentity(_ context.Context, _ *sts.Assu
 func TestBackendSecurityPolicyController_ReconcileOIDC(t *testing.T) {
 	syncFn := internaltesting.NewSyncFnImpl[aigv1a1.AIServiceBackend]()
 	cl := fake.NewClientBuilder().WithScheme(scheme).Build()
-	c := newBackendSecurityPolicyController(cl, fake2.NewClientset(), ctrl.Log, syncFn.Sync)
+	c := NewBackendSecurityPolicyController(cl, fake2.NewClientset(), ctrl.Log, syncFn.Sync)
 	backendSecurityPolicyName := "mybackendSecurityPolicy"
 	namespace := "default"
 
@@ -129,7 +129,7 @@ func TestBackendSecurityPolicyController_ReconcileOIDC(t *testing.T) {
 
 func TestBackendSecurityController_RotateCredentials(t *testing.T) {
 	cl := fake.NewClientBuilder().WithScheme(scheme).Build()
-	c := newBackendSecurityPolicyController(cl, fake2.NewClientset(), ctrl.Log, internaltesting.NewSyncFnImpl[aigv1a1.AIServiceBackend]().Sync)
+	c := NewBackendSecurityPolicyController(cl, fake2.NewClientset(), ctrl.Log, internaltesting.NewSyncFnImpl[aigv1a1.AIServiceBackend]().Sync)
 	backendSecurityPolicyName := "mybackendSecurityPolicy"
 	namespace := "default"
 
diff --git a/internal/extproc/chatcompletion_processor_test.go b/internal/extproc/chatcompletion_processor_test.go
index 649877a..4c035d5 100644
--- a/internal/extproc/chatcompletion_processor_test.go
+++ b/internal/extproc/chatcompletion_processor_test.go
@@ -14,10 +14,12 @@ import (
 
 	corev3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
 	extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
+	typev3 "github.com/envoyproxy/go-control-plane/envoy/type/v3"
 	"github.com/stretchr/testify/require"
 
 	"github.com/envoyproxy/ai-gateway/filterapi"
 	"github.com/envoyproxy/ai-gateway/internal/apischema/openai"
+	"github.com/envoyproxy/ai-gateway/internal/extproc/router"
 	"github.com/envoyproxy/ai-gateway/internal/extproc/translator"
 	"github.com/envoyproxy/ai-gateway/internal/llmcostcel"
 )
@@ -161,6 +163,18 @@ func TestChatCompletion_ProcessRequestBody(t *testing.T) {
 		_, err := p.ProcessRequestBody(t.Context(), &extprocv3.HttpBody{Body: bodyFromModel(t, "some-model")})
 		require.ErrorContains(t, err, "failed to calculate route: test error")
 	})
+	t.Run("router error 404", func(t *testing.T) {
+		headers := map[string]string{":path": "/foo"}
+		rt := mockRouter{t: t, expHeaders: headers, retErr: router.ErrNoMatchingRule}
+		p := &chatCompletionProcessor{config: &processorConfig{router: rt}, requestHeaders: headers, logger: slog.Default()}
+		resp, err := p.ProcessRequestBody(t.Context(), &extprocv3.HttpBody{Body: bodyFromModel(t, "some-model")})
+		require.NoError(t, err)
+		require.NotNil(t, resp)
+		ir := resp.GetImmediateResponse()
+		require.NotNil(t, ir)
+		require.Equal(t, typev3.StatusCode_NotFound, ir.GetStatus().GetCode())
+		require.Equal(t, router.ErrNoMatchingRule.Error(), string(ir.GetBody()))
+	})
 	t.Run("translator not found", func(t *testing.T) {
 		headers := map[string]string{":path": "/foo"}
 		rt := mockRouter{
diff --git a/tests/controller/controller_test.go b/tests/controller/controller_test.go
index 98efdc3..9b9ff76 100644
--- a/tests/controller/controller_test.go
+++ b/tests/controller/controller_test.go
@@ -490,7 +490,125 @@ func TestAIGatewayRouteController(t *testing.T) {
 	})
 }
 
-func TestBackendSecurityPolicyController(t *testing.T) { t.Skip("TODO") }
+func TestBackendSecurityPolicyController(t *testing.T) {
+	c, cfg, k := testsinternal.NewEnvTest(t)
+
+	syncAIServiceBackend := internaltesting.NewSyncFnImpl[aigv1a1.AIServiceBackend]()
+
+	opt := ctrl.Options{Scheme: c.Scheme(), LeaderElection: false, Controller: config.Controller{SkipNameValidation: ptr.To(true)}}
+	mgr, err := ctrl.NewManager(cfg, opt)
+	require.NoError(t, err)
+	require.NoError(t, controller.ApplyIndexing(t.Context(), mgr.GetFieldIndexer().IndexField))
+
+	pc := controller.NewBackendSecurityPolicyController(mgr.GetClient(), k, defaultLogger(), syncAIServiceBackend.Sync)
+	err = ctrl.NewControllerManagedBy(mgr).For(&aigv1a1.BackendSecurityPolicy{}).Complete(pc)
+	require.NoError(t, err)
+
+	go func() {
+		err := mgr.Start(t.Context())
+		require.NoError(t, err)
+	}()
+
+	const backendSecurityPolicyName, backendSecurityPolicyNamespace = "bsp", "default"
+
+	originals := []*aigv1a1.AIServiceBackend{
+		{
+			ObjectMeta: metav1.ObjectMeta{Name: "backend1", Namespace: backendSecurityPolicyNamespace},
+			Spec: aigv1a1.AIServiceBackendSpec{
+				APISchema: defaultSchema,
+				BackendRef: gwapiv1.BackendObjectReference{
+					Name: gwapiv1.ObjectName("mybackend"),
+					Port: ptr.To[gwapiv1.PortNumber](8080),
+				},
+				BackendSecurityPolicyRef: &gwapiv1.LocalObjectReference{
+					Kind:  "BackendSecurityPolicy",
+					Group: "aigateway.envoyproxy.io",
+					Name:  gwapiv1.ObjectName(backendSecurityPolicyName),
+				},
+			},
+		},
+		{
+			ObjectMeta: metav1.ObjectMeta{Name: "backend2", Namespace: backendSecurityPolicyNamespace},
+			Spec: aigv1a1.AIServiceBackendSpec{
+				APISchema: defaultSchema,
+				BackendRef: gwapiv1.BackendObjectReference{
+					Name: gwapiv1.ObjectName("mybackend"),
+					Port: ptr.To[gwapiv1.PortNumber](8080),
+				},
+				BackendSecurityPolicyRef: &gwapiv1.LocalObjectReference{
+					Kind:  "BackendSecurityPolicy",
+					Group: "aigateway.envoyproxy.io",
+					Name:  gwapiv1.ObjectName(backendSecurityPolicyName),
+				},
+			},
+		},
+	}
+	for _, backend := range originals {
+		require.NoError(t, c.Create(t.Context(), backend))
+	}
+
+	t.Run("create security policy", func(t *testing.T) {
+		origin := &aigv1a1.BackendSecurityPolicy{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:      backendSecurityPolicyName,
+				Namespace: backendSecurityPolicyNamespace,
+			},
+			Spec: aigv1a1.BackendSecurityPolicySpec{
+				Type: aigv1a1.BackendSecurityPolicyTypeAPIKey,
+				APIKey: &aigv1a1.BackendSecurityPolicyAPIKey{
+					SecretRef: &gwapiv1.SecretObjectReference{
+						Name: "secret",
+					},
+				},
+			},
+		}
+		require.NoError(t, c.Create(t.Context(), origin))
+		require.Eventually(t, func() bool {
+			println("Eeee")
+			return len(syncAIServiceBackend.GetItems()) == 2
+		}, 5*time.Second, 200*time.Millisecond)
+
+		// Verify that they are the same.
+		backends := syncAIServiceBackend.GetItems()
+		sort.Slice(backends, func(i, j int) bool {
+			backends[i].TypeMeta = metav1.TypeMeta{}
+			backends[j].TypeMeta = metav1.TypeMeta{}
+			return backends[i].Name < backends[j].Name
+		})
+		require.Equal(t, originals, backends)
+	})
+
+	syncAIServiceBackend.Reset()
+	t.Run("update security policy", func(t *testing.T) {
+		origin := &aigv1a1.BackendSecurityPolicy{}
+		require.NoError(t, c.Get(t.Context(), client.ObjectKey{Name: backendSecurityPolicyName, Namespace: backendSecurityPolicyNamespace}, origin))
+		origin.Spec.Type = aigv1a1.BackendSecurityPolicyTypeAWSCredentials
+		origin.Spec.APIKey = nil
+		origin.Spec.AWSCredentials = &aigv1a1.BackendSecurityPolicyAWSCredentials{
+			Region: "us-east-1",
+			CredentialsFile: &aigv1a1.AWSCredentialsFile{
+				SecretRef: &gwapiv1.SecretObjectReference{
+					Name:      "secret",
+					Namespace: ptr.To[gwapiv1.Namespace](backendSecurityPolicyNamespace),
+				},
+			},
+		}
+		require.NoError(t, c.Update(t.Context(), origin))
+
+		require.Eventually(t, func() bool {
+			return len(syncAIServiceBackend.GetItems()) == 2
+		}, 5*time.Second, 200*time.Millisecond)
+
+		// Verify that they are the same.
+		backends := syncAIServiceBackend.GetItems()
+		sort.Slice(backends, func(i, j int) bool {
+			backends[i].TypeMeta = metav1.TypeMeta{}
+			backends[j].TypeMeta = metav1.TypeMeta{}
+			return backends[i].Name < backends[j].Name
+		})
+		require.Equal(t, originals, backends)
+	})
+}
 
 func TestAIServiceBackendController(t *testing.T) {
 	c, cfg, k := testsinternal.NewEnvTest(t)