add ut for sensitive component

Author: Lei Da

component/sensitive.go

@@ -13,7 +13,6 @@ import (
 
 type sensitiveComponentImpl struct {
 	checker rpc.ModerationSvcClient
-	enable  bool
 }
 
 type SensitiveComponent interface {
@@ -23,20 +22,16 @@ type SensitiveComponent interface {
 }
 
 func NewSensitiveComponent(cfg *config.Config) (SensitiveComponent, error) {
-	c := &sensitiveComponentImpl{}
-	c.enable = cfg.SensitiveCheck.Enable
-
-	if c.enable {
-		c.checker = rpc.NewModerationSvcHttpClient(fmt.Sprintf("%s:%d", cfg.Moderation.Host, cfg.Moderation.Port))
+	if !cfg.SensitiveCheck.Enable {
+		return &sensitiveComponentNoOpImpl{}, nil
 	}
+
+	c := &sensitiveComponentImpl{}
+	c.checker = rpc.NewModerationSvcHttpClient(fmt.Sprintf("%s:%d", cfg.Moderation.Host, cfg.Moderation.Port))
 	return c, nil
 }
 
-func (c *sensitiveComponentImpl) CheckText(ctx context.Context, scenario, text string) (bool, error) {
-	if !c.enable {
-		return true, nil
-	}
-
+func (c sensitiveComponentImpl) CheckText(ctx context.Context, scenario, text string) (bool, error) {
 	result, err := c.checker.PassTextCheck(ctx, scenario, text)
 	if err != nil {
 		return false, err
@@ -45,23 +40,15 @@ func (c *sensitiveComponentImpl) CheckText(ctx context.Context, scenario, text s
 	return !result.IsSensitive, nil
 }
 
-func (c *sensitiveComponentImpl) CheckImage(ctx context.Context, scenario, ossBucketName, ossObjectName string) (bool, error) {
-	if !c.enable {
-		return true, nil
-	}
-
+func (c sensitiveComponentImpl) CheckImage(ctx context.Context, scenario, ossBucketName, ossObjectName string) (bool, error) {
 	result, err := c.checker.PassImageCheck(ctx, scenario, ossBucketName, ossObjectName)
 	if err != nil {
 		return false, err
 	}
 	return !result.IsSensitive, nil
 }
 
-func (c *sensitiveComponentImpl) CheckRequestV2(ctx context.Context, req types.SensitiveRequestV2) (bool, error) {
-	if !c.enable {
-		return true, nil
-	}
-
+func (c sensitiveComponentImpl) CheckRequestV2(ctx context.Context, req types.SensitiveRequestV2) (bool, error) {
 	fields := req.GetSensitiveFields()
 	for _, field := range fields {
 		if len(field.Value()) == 0 {
@@ -79,3 +66,21 @@ func (c *sensitiveComponentImpl) CheckRequestV2(ctx context.Context, req types.S
 	}
 	return true, nil
 }
+
+// sensitiveComponentNoOpImpl this implementation provides a "no-op" (no operation) version of the SensitiveComponent interface,
+// where all methods simply return a "not sensitive" result without performing any actual checks.
+type sensitiveComponentNoOpImpl struct {
+}
+
+func (c *sensitiveComponentNoOpImpl) CheckText(ctx context.Context, scenario, text string) (bool, error) {
+	return true, nil
+}
+
+func (c *sensitiveComponentNoOpImpl) CheckImage(ctx context.Context, scenario, ossBucketName, ossObjectName string) (bool, error) {
+	return true, nil
+}
+
+// implements SensitiveComponent
+func (c *sensitiveComponentNoOpImpl) CheckRequestV2(ctx context.Context, req types.SensitiveRequestV2) (bool, error) {
+	return true, nil
+}

component/sensitive_test.go

@@ -0,0 +1,94 @@
+package component
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+	mockrpc "opencsg.com/csghub-server/_mocks/opencsg.com/csghub-server/builder/rpc"
+	mocktypes "opencsg.com/csghub-server/_mocks/opencsg.com/csghub-server/common/types"
+	"opencsg.com/csghub-server/builder/rpc"
+	"opencsg.com/csghub-server/builder/sensitive"
+	"opencsg.com/csghub-server/common/config"
+	"opencsg.com/csghub-server/common/types"
+)
+
+func TestSensitiveComponent_CheckText(t *testing.T) {
+	mockModeration := mockrpc.NewMockModerationSvcClient(t)
+	mockModeration.EXPECT().PassTextCheck(mock.Anything, mock.Anything, mock.Anything).Return(&rpc.CheckResult{
+		IsSensitive: false,
+	}, nil)
+	comp := &sensitiveComponentImpl{
+		checker: mockModeration,
+	}
+
+	success, err := comp.CheckText(context.TODO(), string(sensitive.ScenarioChatDetection), "test")
+	require.Nil(t, err)
+	require.True(t, success)
+}
+
+func TestSensitiveComponent_CheckImage(t *testing.T) {
+	mockModeration := mockrpc.NewMockModerationSvcClient(t)
+	mockModeration.EXPECT().PassImageCheck(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(&rpc.CheckResult{
+		IsSensitive: false,
+	}, nil)
+	comp := &sensitiveComponentImpl{
+		checker: mockModeration,
+	}
+
+	success, err := comp.CheckImage(context.TODO(), string(sensitive.ScenarioChatDetection), "ossBucketName", "ossObjectName")
+	require.Nil(t, err)
+	require.True(t, success)
+}
+
+func TestSensitiveComponent_CheckRequestV2(t *testing.T) {
+	mockModeration := mockrpc.NewMockModerationSvcClient(t)
+	mockModeration.EXPECT().PassTextCheck(mock.Anything, mock.Anything, mock.Anything).Return(&rpc.CheckResult{
+		IsSensitive: false,
+	}, nil).Twice()
+	comp := &sensitiveComponentImpl{
+		checker: mockModeration,
+	}
+	mockRequest := mocktypes.NewMockSensitiveRequestV2(t)
+	mockRequest.EXPECT().GetSensitiveFields().Return([]types.SensitiveField{
+		{
+			Name: "chat",
+			Value: func() string {
+				return "chat1"
+			},
+			Scenario: string(sensitive.ScenarioChatDetection),
+		},
+		{
+			Name: "comment",
+			Value: func() string {
+				return "comment1"
+			},
+			Scenario: string(sensitive.ScenarioCommentDetection),
+		},
+	})
+	success, err := comp.CheckRequestV2(context.TODO(), mockRequest)
+	require.Nil(t, err)
+	require.True(t, success)
+}
+
+func TestSensitiveComponent_NoOpImpl(t *testing.T) {
+	cfg := &config.Config{}
+	cfg.SensitiveCheck.Enable = false
+	c, err := NewSensitiveComponent(cfg)
+	require.Nil(t, err)
+
+	success, err := c.CheckText(context.Background(), string(sensitive.ScenarioChatDetection), "test")
+	require.Nil(t, err)
+	require.True(t, success)
+
+	success, err = c.CheckImage(context.Background(), string(sensitive.ScenarioChatDetection), "test", "test")
+	require.Nil(t, err)
+	require.True(t, success)
+
+	mockRequest := mocktypes.NewMockSensitiveRequestV2(t)
+
+	success, err = c.CheckRequestV2(context.Background(), mockRequest)
+	require.Nil(t, err)
+	require.True(t, success)
+}