github.com/kubeshop/testkube@v1.17.23/pkg/triggers/lease_backend_test.go (about)

     1  package triggers
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/stretchr/testify/assert"
     9  	"go.mongodb.org/mongo-driver/bson"
    10  	"go.mongodb.org/mongo-driver/mongo/integration/mtest"
    11  )
    12  
    13  func TestMongoLeaseBackend_TryAcquire(t *testing.T) {
    14  	t.Parallel()
    15  
    16  	ctx := context.Background()
    17  
    18  	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
    19  	defer mt.Close()
    20  
    21  	mt.Run("acquire existing lease", func(mt *mtest.T) {
    22  		mt.Parallel()
    23  
    24  		leaseBackend := NewMongoLeaseBackend(mt.DB)
    25  
    26  		testID := "test-host-1"
    27  		testClusterID := "testkube_api"
    28  		expectedLease := Lease{
    29  			Identifier: testID,
    30  			ClusterID:  testClusterID,
    31  			AcquiredAt: time.Now(),
    32  			RenewedAt:  time.Now(),
    33  		}
    34  		expectedMongoLease := MongoLease{
    35  			_id:   newLeaseMongoID(testClusterID),
    36  			Lease: expectedLease,
    37  		}
    38  
    39  		bsonD := cycleBSON(&expectedMongoLease)
    40  
    41  		mt.AddMockResponses(mtest.CreateCursorResponse(
    42  			1,
    43  			"testkube.lease",
    44  			mtest.FirstBatch,
    45  			bsonD,
    46  		))
    47  
    48  		leased, err := leaseBackend.TryAcquire(ctx, testID, testClusterID)
    49  
    50  		assert.True(t, leased, "should acquire lease")
    51  		assert.NoError(t, err)
    52  	})
    53  
    54  	mt.Run("renew existing lease", func(mt *mtest.T) {
    55  		mt.Parallel()
    56  
    57  		leaseBackend := NewMongoLeaseBackend(mt.DB)
    58  
    59  		testID := "test-host-1"
    60  		testClusterID := "testkube_api"
    61  		acquiredAt := time.Now().Add(-1 * time.Hour)
    62  		expectedLease1 := Lease{
    63  			Identifier: testID,
    64  			ClusterID:  testClusterID,
    65  			AcquiredAt: acquiredAt,
    66  			RenewedAt:  time.Now().Add(-1 * time.Hour),
    67  		}
    68  		expectedMongoLease1 := MongoLease{
    69  			_id:   newLeaseMongoID(testClusterID),
    70  			Lease: expectedLease1,
    71  		}
    72  		expectedLease2 := Lease{
    73  			Identifier: testID,
    74  			ClusterID:  testClusterID,
    75  			AcquiredAt: acquiredAt,
    76  			RenewedAt:  time.Now(),
    77  		}
    78  		expectedMongoLease2 := MongoLease{
    79  			_id:   newLeaseMongoID(testClusterID),
    80  			Lease: expectedLease2,
    81  		}
    82  
    83  		bsonD1 := cycleBSON(&expectedMongoLease1)
    84  		bsonD2 := cycleBSON(&expectedMongoLease2)
    85  
    86  		mt.AddMockResponses(mtest.CreateCursorResponse(
    87  			1,
    88  			"testkube.lease",
    89  			mtest.FirstBatch,
    90  			bsonD1,
    91  		))
    92  		mt.AddMockResponses(mtest.CreateCursorResponse(
    93  			0,
    94  			"testkube.lease",
    95  			mtest.FirstBatch,
    96  		))
    97  		mt.AddMockResponses(bson.D{
    98  			{Key: "ok", Value: 1},
    99  			{Key: "value", Value: bsonD2},
   100  		})
   101  
   102  		leased, err := leaseBackend.TryAcquire(ctx, testID, testClusterID)
   103  
   104  		assert.True(t, leased, "should acquire lease")
   105  		assert.NoError(t, err)
   106  	})
   107  
   108  	mt.Run("not acquire if other instance is holding non-expired", func(mt *mtest.T) {
   109  		mt.Parallel()
   110  
   111  		leaseBackend := NewMongoLeaseBackend(mt.DB)
   112  
   113  		testClusterID := "testkube_api"
   114  		acquiredAt := time.Now().Add(-1 * time.Hour)
   115  		expectedLease := Lease{
   116  			Identifier: "test-id-2",
   117  			ClusterID:  testClusterID,
   118  			AcquiredAt: acquiredAt,
   119  			RenewedAt:  time.Now().Add(-5 * time.Second),
   120  		}
   121  		expectedMongoLease := MongoLease{
   122  			_id:   newLeaseMongoID(testClusterID),
   123  			Lease: expectedLease,
   124  		}
   125  
   126  		bsonD := cycleBSON(&expectedMongoLease)
   127  
   128  		mt.AddMockResponses(mtest.CreateCursorResponse(
   129  			1,
   130  			"testkube.lease",
   131  			mtest.FirstBatch,
   132  			bsonD,
   133  		))
   134  
   135  		leased, err := leaseBackend.TryAcquire(ctx, "test-host-1", testClusterID)
   136  
   137  		assert.False(t, leased, "should not acquire lease")
   138  		assert.NoError(t, err)
   139  	})
   140  
   141  	mt.Run("acquire lease from other instance if lease is expired", func(mt *mtest.T) {
   142  		mt.Parallel()
   143  
   144  		leaseBackend := NewMongoLeaseBackend(mt.DB)
   145  
   146  		testClusterID := "testkube_api"
   147  		acquiredAt := time.Now().Add(-1 * time.Hour)
   148  		expectedLease1 := Lease{
   149  			Identifier: "test-host-2",
   150  			ClusterID:  testClusterID,
   151  			AcquiredAt: acquiredAt,
   152  			RenewedAt:  time.Now().Add(-1 * time.Hour),
   153  		}
   154  		expectedMongoLease1 := MongoLease{
   155  			_id:   newLeaseMongoID(testClusterID),
   156  			Lease: expectedLease1,
   157  		}
   158  		expectedLease2 := Lease{
   159  			Identifier: "test-host-1",
   160  			ClusterID:  testClusterID,
   161  			AcquiredAt: acquiredAt,
   162  			RenewedAt:  time.Now(),
   163  		}
   164  		expectedMongoLease2 := MongoLease{
   165  			_id:   newLeaseMongoID(testClusterID),
   166  			Lease: expectedLease2,
   167  		}
   168  
   169  		bsonD1 := cycleBSON(&expectedMongoLease1)
   170  		bsonD2 := cycleBSON(&expectedMongoLease2)
   171  
   172  		mt.AddMockResponses(mtest.CreateCursorResponse(
   173  			1,
   174  			"testkube.lease",
   175  			mtest.FirstBatch,
   176  			bsonD1,
   177  		))
   178  		mt.AddMockResponses(mtest.CreateCursorResponse(
   179  			0,
   180  			"testkube.lease",
   181  			mtest.FirstBatch,
   182  		))
   183  		mt.AddMockResponses(bson.D{
   184  			{Key: "ok", Value: 1},
   185  			{Key: "value", Value: bsonD2},
   186  		})
   187  
   188  		leased, err := leaseBackend.TryAcquire(ctx, "test-host-1", testClusterID)
   189  
   190  		assert.True(t, leased, "should acquire lease")
   191  		assert.NoError(t, err)
   192  	})
   193  }
   194  
   195  func cycleBSON(data any) bson.D {
   196  	bsonData, _ := bson.Marshal(data)
   197  	var bsonD bson.D
   198  	_ = bson.Unmarshal(bsonData, &bsonD)
   199  
   200  	return bsonD
   201  }