github.com/castai/kvisor@v1.7.1-0.20240516114728-b3572a2607b5/cmd/controller/state/castai_controller_test.go (about)

     1  package state
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"log/slog"
     7  	"testing"
     8  	"time"
     9  
    10  	castaipb "github.com/castai/kvisor/api/v1/runtime"
    11  	"github.com/castai/kvisor/cmd/controller/kube"
    12  	"github.com/castai/kvisor/pkg/castai"
    13  	"github.com/castai/kvisor/pkg/logging"
    14  	"github.com/stretchr/testify/require"
    15  	"google.golang.org/grpc"
    16  	"k8s.io/client-go/kubernetes/fake"
    17  )
    18  
    19  func TestCastaiController(t *testing.T) {
    20  	log := logging.New(&logging.Config{
    21  		Level: slog.LevelDebug,
    22  	})
    23  	k8sClient := fake.NewSimpleClientset()
    24  	kubeClient := kube.NewClient(log, "agent", "ns", kube.Version{}, k8sClient)
    25  	ctx, cancel := context.WithCancel(context.Background())
    26  	defer cancel()
    27  
    28  	t.Run("stop on initial config failure", func(t *testing.T) {
    29  		r := require.New(t)
    30  		client := &testGrpcClient{
    31  			getConfigurationResponse: func() (*castaipb.GetConfigurationResponse, error) {
    32  				return nil, errors.New("ups")
    33  			},
    34  		}
    35  		ctrl := newTestCastaiController(log, kubeClient, client)
    36  		ctrl.remoteConfigInitialSyncTimeout = 50 * time.Millisecond
    37  		ctrl.remoteConfigRetryWaitDuration = 10 * time.Millisecond
    38  		err := ctrl.Run(ctx)
    39  		r.ErrorIs(err, context.DeadlineExceeded)
    40  	})
    41  
    42  	t.Run("stop on config loop failure after max retries", func(t *testing.T) {
    43  		r := require.New(t)
    44  		calls := 0
    45  		client := &testGrpcClient{
    46  			getConfigurationResponse: func() (*castaipb.GetConfigurationResponse, error) {
    47  				calls++
    48  				// Do not fail to first call to pass initial config sync.
    49  				if calls == 1 {
    50  					return &castaipb.GetConfigurationResponse{}, nil
    51  				}
    52  				return nil, errors.New("ups")
    53  			},
    54  		}
    55  		ctrl := newTestCastaiController(log, kubeClient, client)
    56  		ctrl.remoteConfigInitialSyncTimeout = 50 * time.Millisecond
    57  		ctrl.remoteConfigRetryWaitDuration = 10 * time.Millisecond
    58  		ctrl.removeConfigMaxFailures = 3
    59  		err := ctrl.Run(ctx)
    60  		r.ErrorContains(err, "remote config fetch errors reached")
    61  	})
    62  }
    63  
    64  func newTestCastaiController(log *logging.Logger, kubeClient *kube.Client, client *testGrpcClient) *CastaiController {
    65  	castaiClient := &castai.Client{
    66  		GRPC: client,
    67  	}
    68  	cfg := CastaiConfig{
    69  		RemoteConfigSyncDuration: 10 * time.Millisecond,
    70  	}
    71  	return NewCastaiController(log, cfg, &castaipb.ControllerConfig{}, kubeClient, castaiClient)
    72  }
    73  
    74  type testGrpcClient struct {
    75  	getConfigurationResponse func() (*castaipb.GetConfigurationResponse, error)
    76  }
    77  
    78  func (t testGrpcClient) KubeBenchReportIngest(ctx context.Context, in *castaipb.KubeBenchReport, opts ...grpc.CallOption) (*castaipb.KubeBenchReportIngestResponse, error) {
    79  	return nil, nil
    80  }
    81  
    82  func (t testGrpcClient) KubeLinterReportIngest(ctx context.Context, in *castaipb.KubeLinterReport, opts ...grpc.CallOption) (*castaipb.KubeLinterReportIngestResponse, error) {
    83  	return nil, nil
    84  }
    85  
    86  func (t testGrpcClient) ImageMetadataIngest(ctx context.Context, in *castaipb.ImageMetadata, opts ...grpc.CallOption) (*castaipb.ImageMetadataIngestResponse, error) {
    87  	return nil, nil
    88  }
    89  
    90  func (t testGrpcClient) GetSyncState(ctx context.Context, in *castaipb.GetSyncStateRequest, opts ...grpc.CallOption) (*castaipb.GetSyncStateResponse, error) {
    91  	return nil, nil
    92  }
    93  
    94  func (t testGrpcClient) UpdateSyncState(ctx context.Context, in *castaipb.UpdateSyncStateRequest, opts ...grpc.CallOption) (*castaipb.UpdateSyncStateResponse, error) {
    95  	return nil, nil
    96  }
    97  
    98  func (t testGrpcClient) GetConfiguration(ctx context.Context, in *castaipb.GetConfigurationRequest, opts ...grpc.CallOption) (*castaipb.GetConfigurationResponse, error) {
    99  	if t.getConfigurationResponse != nil {
   100  		return t.getConfigurationResponse()
   101  	}
   102  	return &castaipb.GetConfigurationResponse{}, nil
   103  }
   104  
   105  func (t testGrpcClient) EventsWriteStream(ctx context.Context, opts ...grpc.CallOption) (castaipb.RuntimeSecurityAgentAPI_EventsWriteStreamClient, error) {
   106  	return nil, nil
   107  }
   108  
   109  func (t testGrpcClient) LogsWriteStream(ctx context.Context, opts ...grpc.CallOption) (castaipb.RuntimeSecurityAgentAPI_LogsWriteStreamClient, error) {
   110  	return nil, nil
   111  }
   112  
   113  func (t testGrpcClient) ContainerStatsWriteStream(ctx context.Context, opts ...grpc.CallOption) (castaipb.RuntimeSecurityAgentAPI_ContainerStatsWriteStreamClient, error) {
   114  	return nil, nil
   115  }
   116  
   117  func (testGrpcClient) KubernetesDeltaIngest(ctx context.Context, opts ...grpc.CallOption) (castaipb.RuntimeSecurityAgentAPI_KubernetesDeltaIngestClient, error) {
   118  	return nil, nil
   119  }