From d42f67e0536af630929d524d57d345bbac1f9ffa Mon Sep 17 00:00:00 2001
From: Ekaterina Lebedeva <ekaterina.lebedeva@yadro.com>
Date: Sat, 22 Mar 2025 16:40:31 +0300
Subject: [PATCH] [#1656] qos: Add tests for `AdjustOutgoingIOTag` Interceptors

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

diff --git a/internal/qos/grpc_test.go b/internal/qos/grpc_test.go
index c5b4cd162..215329c86 100644
--- a/internal/qos/grpc_test.go
+++ b/internal/qos/grpc_test.go
@@ -3,6 +3,7 @@ package qos_test
 import (
 	"context"
 	"errors"
+	"fmt"
 	"testing"
 
 	"git.frostfs.info/TrueCloudLab/frostfs-node/internal/qos"
@@ -20,7 +21,9 @@ const (
 var (
 	errTest         = errors.New("mock")
 	errWrongTag     = errors.New("wrong tag")
+	errNoTag        = errors.New("failed to get tag from context")
 	errResExhausted = new(apistatus.ResourceExhausted)
+	tags            = []qos.IOTag{qos.IOTagBackground, qos.IOTagWritecache, qos.IOTagPolicer, qos.IOTagTreeSync}
 )
 
 type mockGRPCServerStream struct {
@@ -130,3 +133,82 @@ func TestSetCriticalIOTagUnaryServerInterceptor_Pass(t *testing.T) {
 	_, err := interceptor(context.Background(), nil, nil, handler)
 	require.NoError(t, err)
 }
+
+func TestAdjustOutgoingIOTagUnaryClientInterceptor(t *testing.T) {
+	interceptor := qos.NewAdjustOutgoingIOTagUnaryClientInterceptor()
+
+	// check context with no value
+	called := false
+	invoker := func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, opts ...grpc.CallOption) error {
+		called = true
+		if _, ok := tagging.IOTagFromContext(ctx); ok {
+			return fmt.Errorf("%v: expected no IO tags", errWrongTag)
+		}
+		return nil
+	}
+	require.NoError(t, interceptor(context.Background(), "", nil, nil, nil, invoker, nil))
+	require.True(t, called)
+
+	// check context for internal tag
+	targetTag := qos.IOTagInternal.String()
+	invoker = func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, opts ...grpc.CallOption) error {
+		raw, ok := tagging.IOTagFromContext(ctx)
+		if !ok {
+			return errNoTag
+		}
+		if raw != targetTag {
+			return errWrongTag
+		}
+		return nil
+	}
+	for _, tag := range tags {
+		ctx := tagging.ContextWithIOTag(context.Background(), tag.String())
+		require.NoError(t, interceptor(ctx, "", nil, nil, nil, invoker, nil))
+	}
+
+	// check context for client tag
+	ctx := tagging.ContextWithIOTag(context.Background(), "")
+	targetTag = qos.IOTagClient.String()
+	require.NoError(t, interceptor(ctx, "", nil, nil, nil, invoker, nil))
+}
+
+func TestAdjustOutgoingIOTagStreamClientInterceptor(t *testing.T) {
+	interceptor := qos.NewAdjustOutgoingIOTagStreamClientInterceptor()
+
+	// check context with no value
+	called := false
+	streamer := func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) {
+		called = true
+		if _, ok := tagging.IOTagFromContext(ctx); ok {
+			return nil, fmt.Errorf("%v: expected no IO tags", errWrongTag)
+		}
+		return nil, nil
+	}
+	_, err := interceptor(context.Background(), nil, nil, "", streamer, nil)
+	require.True(t, called)
+	require.NoError(t, err)
+
+	// check context for internal tag
+	targetTag := qos.IOTagInternal.String()
+	streamer = func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) {
+		raw, ok := tagging.IOTagFromContext(ctx)
+		if !ok {
+			return nil, errNoTag
+		}
+		if raw != targetTag {
+			return nil, errWrongTag
+		}
+		return nil, nil
+	}
+	for _, tag := range tags {
+		ctx := tagging.ContextWithIOTag(context.Background(), tag.String())
+		_, err := interceptor(ctx, nil, nil, "", streamer, nil)
+		require.NoError(t, err)
+	}
+
+	// check context for client tag
+	ctx := tagging.ContextWithIOTag(context.Background(), "")
+	targetTag = qos.IOTagClient.String()
+	_, err = interceptor(ctx, nil, nil, "", streamer, nil)
+	require.NoError(t, err)
+}