github.com/influxdata/influxdb/v2@v2.7.6/replications/transport/middleware_kv_test.go (about)

     1  package transport
     2  
     3  import (
     4  	"context"
     5  	"encoding/binary"
     6  	"testing"
     7  
     8  	"github.com/golang/mock/gomock"
     9  	"github.com/influxdata/influxdb/v2"
    10  	"github.com/influxdata/influxdb/v2/inmem"
    11  	"github.com/influxdata/influxdb/v2/kit/platform"
    12  	"github.com/influxdata/influxdb/v2/kv"
    13  	"github.com/influxdata/influxdb/v2/kv/migration"
    14  	"github.com/influxdata/influxdb/v2/kv/migration/all"
    15  	"github.com/influxdata/influxdb/v2/replications/mock"
    16  	"github.com/stretchr/testify/require"
    17  	"go.uber.org/zap"
    18  )
    19  
    20  func getCount(kvStore kv.Store, orgID platform.ID) (uint64, error) {
    21  	var count uint64
    22  	if err := kvStore.Update(context.Background(), func(tx kv.Tx) error {
    23  		encodedID, err := orgID.Encode()
    24  		if err != nil {
    25  			return err
    26  		}
    27  		bucket, err := tx.Bucket([]byte("replicationsv2"))
    28  		if err != nil {
    29  			return err
    30  		}
    31  		c, err := bucket.Get(encodedID)
    32  		if err != nil {
    33  			return err
    34  		}
    35  
    36  		count = binary.BigEndian.Uint64(c)
    37  		return nil
    38  	}); err != nil {
    39  		return 0, err
    40  	}
    41  
    42  	return count, nil
    43  }
    44  
    45  func TestReplicationCreateKVUpdate(t *testing.T) {
    46  	kvStore := inmem.NewKVStore()
    47  	gmock := gomock.NewController(t)
    48  	defer gmock.Finish()
    49  	mockRemote := mock.NewMockReplicationService(gmock)
    50  	telemetry := newTelemetryCollectingService(kvStore, mockRemote)
    51  
    52  	kvMigrator, err := migration.NewMigrator(
    53  		zap.L(),
    54  		kvStore,
    55  		all.Migrations[:]...,
    56  	)
    57  	require.NoError(t, err)
    58  	require.NoError(t, kvMigrator.Up(context.Background()))
    59  
    60  	ctx := context.Background()
    61  	req := influxdb.CreateReplicationRequest{
    62  		OrgID:          platform.ID(1),
    63  		Name:           "test1",
    64  		RemoteBucketID: platform.ID(11),
    65  		LocalBucketID:  platform.ID(22),
    66  	}
    67  
    68  	replication := influxdb.Replication{
    69  		OrgID: platform.ID(1),
    70  	}
    71  	replications := influxdb.Replications{
    72  		Replications: []influxdb.Replication{replication},
    73  	}
    74  
    75  	mockRemote.EXPECT().CreateReplication(ctx, req).Return(&replication, nil).Times(1)
    76  	mockRemote.EXPECT().ListReplications(ctx, influxdb.ReplicationListFilter{OrgID: req.OrgID}).Return(&replications, nil).Times(1)
    77  
    78  	repl, err := telemetry.CreateReplication(ctx, req)
    79  	require.NoError(t, err)
    80  
    81  	count, err := getCount(kvStore, repl.OrgID)
    82  	require.NoError(t, err)
    83  	require.Equal(t, int64(1), int64(count))
    84  }
    85  
    86  func TestReplicationDeleteKVUpdate(t *testing.T) {
    87  	kvStore := inmem.NewKVStore()
    88  	gmock := gomock.NewController(t)
    89  	defer gmock.Finish()
    90  	mockRemote := mock.NewMockReplicationService(gmock)
    91  	telemetry := newTelemetryCollectingService(kvStore, mockRemote)
    92  
    93  	ctx := context.Background()
    94  
    95  	kvMigrator, err := migration.NewMigrator(
    96  		zap.L(),
    97  		kvStore,
    98  		all.Migrations[:]...,
    99  	)
   100  	require.NoError(t, err)
   101  	require.NoError(t, kvMigrator.Up(ctx))
   102  
   103  	req := influxdb.CreateReplicationRequest{
   104  		OrgID:          platform.ID(1),
   105  		Name:           "test1",
   106  		RemoteBucketID: platform.ID(11),
   107  		LocalBucketID:  platform.ID(22),
   108  	}
   109  	req2 := req
   110  	req2.Name = "test2"
   111  
   112  	replication1 := influxdb.Replication{
   113  		ID:    platform.ID(1),
   114  		OrgID: platform.ID(1),
   115  	}
   116  	replication2 := replication1
   117  	replication2.ID = platform.ID(2)
   118  
   119  	remoteConnectionsPreDelete := influxdb.Replications{
   120  		Replications: []influxdb.Replication{replication1, replication2},
   121  	}
   122  
   123  	remoteConnectionsPostDelete := influxdb.Replications{
   124  		Replications: []influxdb.Replication{replication1},
   125  	}
   126  
   127  	mockRemote.EXPECT().CreateReplication(ctx, req).Return(&replication1, nil).Times(1)
   128  	mockRemote.EXPECT().CreateReplication(ctx, req2).Return(&replication2, nil).Times(1)
   129  	mockRemote.EXPECT().ListReplications(ctx, influxdb.ReplicationListFilter{OrgID: req.OrgID}).Return(&remoteConnectionsPreDelete, nil).Times(2)
   130  
   131  	mockRemote.EXPECT().DeleteReplication(ctx, replication1.ID).Return(nil).Times(1)
   132  	mockRemote.EXPECT().GetReplication(ctx, replication1.ID).Return(&replication1, nil).Times(1)
   133  	mockRemote.EXPECT().ListReplications(ctx, influxdb.ReplicationListFilter{OrgID: req.OrgID}).Return(&remoteConnectionsPostDelete, nil).Times(1)
   134  
   135  	_, err = telemetry.CreateReplication(ctx, req)
   136  	require.NoError(t, err)
   137  
   138  	repl, err := telemetry.CreateReplication(ctx, req2)
   139  	require.NoError(t, err)
   140  
   141  	err = telemetry.DeleteReplication(ctx, replication1.ID)
   142  	require.NoError(t, err)
   143  
   144  	count, err := getCount(kvStore, repl.OrgID)
   145  	require.NoError(t, err)
   146  	require.Equal(t, int64(1), int64(count))
   147  }