From 115aae7c34bd85f11a866d62ebe453c393f55c70 Mon Sep 17 00:00:00 2001
From: Ekaterina Lebedeva <ekaterina.lebedeva@yadro.com>
Date: Sat, 22 Mar 2025 01:54:06 +0300
Subject: [PATCH] [#1656] qos: Add tests for MaxActiveRPCLimiter Interceptors

Change-Id: Ib65890ae5aec34c34e15d4ec1f05952f74f1ad26
Signed-off-by: Ekaterina Lebedeva <ekaterina.lebedeva@yadro.com>
---
 internal/qos/grpc_test.go | 121 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 121 insertions(+)
 create mode 100644 internal/qos/grpc_test.go

diff --git a/internal/qos/grpc_test.go b/internal/qos/grpc_test.go
new file mode 100644
index 000000000..d6e2a689c
--- /dev/null
+++ b/internal/qos/grpc_test.go
@@ -0,0 +1,121 @@
+package qos_test
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"git.frostfs.info/TrueCloudLab/frostfs-node/internal/qos"
+	"git.frostfs.info/TrueCloudLab/frostfs-qos/limiting"
+	"git.frostfs.info/TrueCloudLab/frostfs-qos/tagging"
+	apistatus "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/client/status"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
+)
+
+const (
+	okKey = "ok"
+)
+
+var (
+	errTest         = errors.New("mock")
+	errResExhausted *apistatus.ResourceExhausted
+)
+
+type mockGRPCServerStream struct {
+	grpc.ServerStream
+
+	ctx context.Context
+}
+
+func (m *mockGRPCServerStream) Context() context.Context {
+	return m.ctx
+}
+
+type limiter struct {
+	acquired bool
+	released bool
+}
+
+func (l *limiter) Acquire(key string) (limiting.ReleaseFunc, bool) {
+	l.acquired = true
+	if key != okKey {
+		return nil, false
+	}
+	return func() { l.released = true }, true
+}
+
+func unaryMaxActiveRPCLimiter(ctx context.Context, lim *limiter, methodName string) error {
+	interceptor := qos.NewMaxActiveRPCLimiterUnaryServerInterceptor(func() limiting.Limiter { return lim })
+	handler := func(ctx context.Context, req any) (any, error) {
+		return nil, errTest
+	}
+	_, err := interceptor(ctx, nil, &grpc.UnaryServerInfo{FullMethod: methodName}, handler)
+	return err
+}
+
+func streamMaxActiveRPCLimiter(ctx context.Context, lim *limiter, methodName string) error {
+	interceptor := qos.NewMaxActiveRPCLimiterStreamServerInterceptor(func() limiting.Limiter { return lim })
+	handler := func(srv any, stream grpc.ServerStream) error {
+		return errTest
+	}
+	err := interceptor(nil, &mockGRPCServerStream{ctx: ctx}, &grpc.StreamServerInfo{
+		FullMethod: methodName,
+	}, handler)
+	return err
+}
+
+func Test_MaxActiveRPCLimiter(t *testing.T) {
+	// UnaryServerInterceptor
+	t.Run("unary fail", func(t *testing.T) {
+		var lim limiter
+
+		err := unaryMaxActiveRPCLimiter(context.Background(), &lim, "")
+		require.ErrorAs(t, err, &errResExhausted)
+		require.True(t, lim.acquired)
+		require.False(t, lim.released)
+	})
+	t.Run("unary pass critical", func(t *testing.T) {
+		var lim limiter
+		ctx := tagging.ContextWithIOTag(context.Background(), qos.IOTagCritical.String())
+
+		err := unaryMaxActiveRPCLimiter(ctx, &lim, "")
+		require.ErrorIs(t, err, errTest)
+		require.False(t, lim.acquired)
+		require.False(t, lim.released)
+	})
+	t.Run("unary pass", func(t *testing.T) {
+		var lim limiter
+
+		err := unaryMaxActiveRPCLimiter(context.Background(), &lim, okKey)
+		require.ErrorIs(t, err, errTest)
+		require.True(t, lim.acquired)
+		require.True(t, lim.released)
+	})
+	// StreamServerInterceptor
+	t.Run("stream fail", func(t *testing.T) {
+		var lim limiter
+
+		err := streamMaxActiveRPCLimiter(context.Background(), &lim, "")
+		require.ErrorAs(t, err, &errResExhausted)
+		require.True(t, lim.acquired)
+		require.False(t, lim.released)
+	})
+	t.Run("stream pass critical", func(t *testing.T) {
+		var lim limiter
+		ctx := tagging.ContextWithIOTag(context.Background(), qos.IOTagCritical.String())
+
+		err := streamMaxActiveRPCLimiter(ctx, &lim, "")
+		require.ErrorIs(t, err, errTest)
+		require.False(t, lim.acquired)
+		require.False(t, lim.released)
+	})
+	t.Run("stream pass", func(t *testing.T) {
+		var lim limiter
+
+		err := streamMaxActiveRPCLimiter(context.Background(), &lim, okKey)
+		require.ErrorIs(t, err, errTest)
+		require.True(t, lim.acquired)
+		require.True(t, lim.released)
+	})
+}