package server import ( "context" "fmt" "sync" "testing" "time" "d7y.io/dragonfly/v2/manager/lease" "d7y.io/dragonfly/v2/manager/server/mocks" "d7y.io/dragonfly/v2/manager/store/client" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/suite" ) type LeaseTestSuite struct { suite.Suite lessor lease.Lessor } func (suite *LeaseTestSuite) TestGrantSuccess() { assert := assert.New(suite.T()) randStr := suite.randPrefix() id1, err := suite.lessor.Grant(context.TODO(), randStr+"key1", randStr+"value1", 10) assert.Nil(err) assert.NotEmpty(id1) id2, err := suite.lessor.Grant(context.TODO(), randStr+"key2", randStr+"value2", 20) assert.Nil(err) assert.NotEmpty(id2) if len(id1) > 0 { suite.lessor.Revoke(context.TODO(), id1) } if len(id2) > 0 { suite.lessor.Revoke(context.TODO(), id2) } } func (suite *LeaseTestSuite) TestGrantFailed() { assert := assert.New(suite.T()) randStr := suite.randPrefix() id1, err := suite.lessor.Grant(context.TODO(), randStr+"key1", randStr+"value1", 10) assert.Nil(err) assert.NotEmpty(id1) id2, err := suite.lessor.Grant(context.TODO(), randStr+"key1", randStr+"value1", 20) assert.NotNil(err) assert.Empty(id2) time.Sleep(time.Duration(5) * time.Second) id3, err := suite.lessor.Grant(context.TODO(), randStr+"key1", randStr+"value1", 20) assert.NotNil(err) assert.Empty(id3) time.Sleep(time.Duration(6) * time.Second) id4, err := suite.lessor.Grant(context.TODO(), randStr+"key1", randStr+"value1", 20) assert.Nil(err) assert.NotEmpty(id4) if len(id1) > 0 { suite.lessor.Revoke(context.TODO(), id1) } if len(id2) > 0 { suite.lessor.Revoke(context.TODO(), id2) } if len(id3) > 0 { suite.lessor.Revoke(context.TODO(), id3) } if len(id4) > 0 { suite.lessor.Revoke(context.TODO(), id4) } } func (suite *LeaseTestSuite) TestKeepAliveSuccess() { assert := assert.New(suite.T()) randStr := suite.randPrefix() id1, err := suite.lessor.Grant(context.TODO(), randStr+"key1", randStr+"value1", 10) assert.Nil(err) assert.NotEmpty(id1) if len(id1) > 0 { ch, err := suite.lessor.KeepAlive(context.TODO(), id1) assert.Nil(err) assert.NotNil(ch) } time.Sleep(time.Duration(15) * time.Second) id2, err := suite.lessor.Grant(context.TODO(), randStr+"key1", randStr+"value1", 10) assert.NotNil(err) assert.Empty(id2) if len(id1) > 0 { suite.lessor.Revoke(context.TODO(), id1) } if len(id2) > 0 { suite.lessor.Revoke(context.TODO(), id2) } } func (suite *LeaseTestSuite) TestKeepAliveTimeout() { assert := assert.New(suite.T()) randStr := suite.randPrefix() id1, err := suite.lessor.Grant(context.TODO(), randStr+"key1", randStr+"value1", 10) assert.Nil(err) assert.NotEmpty(id1) if len(id1) > 0 { ch, err := suite.lessor.KeepAlive(context.TODO(), id1) assert.Nil(err) assert.NotNil(ch) } time.Sleep(time.Duration(15) * time.Second) id2, err := suite.lessor.Grant(context.TODO(), randStr+"key1", randStr+"value1", 10) assert.NotNil(err) assert.Empty(id2) suite.lessor.Close() defer func() { suite.lessor = nil }() time.Sleep(time.Duration(15) * time.Second) id3, err := suite.lessor.Grant(context.TODO(), randStr+"key1", randStr+"value1", 10) assert.Nil(err) assert.NotEmpty(id3) if len(id1) > 0 { suite.lessor.Revoke(context.TODO(), id1) } if len(id2) > 0 { suite.lessor.Revoke(context.TODO(), id2) } if len(id3) > 0 { suite.lessor.Revoke(context.TODO(), id2) } } func (suite *LeaseTestSuite) TestKeepAliveMeetRevoke() { assert := assert.New(suite.T()) randStr := suite.randPrefix() id1, err := suite.lessor.Grant(context.TODO(), randStr+"key1", randStr+"value1", 10) assert.Nil(err) assert.NotEmpty(id1) ch, err := suite.lessor.KeepAlive(context.TODO(), id1) assert.Nil(err) assert.NotNil(ch) var wg sync.WaitGroup wg.Add(1) go func(ch chan struct{}) { defer wg.Done() for { select { case <-ch: return } } }(ch) time.Sleep(time.Duration(15) * time.Second) suite.lessor.Revoke(context.TODO(), id1) wg.Wait() } func (suite *LeaseTestSuite) randPrefix() string { return fmt.Sprintf("%d_", time.Now().Unix()) } func (suite *LeaseTestSuite) SetupTest() { assert := assert.New(suite.T()) cfg := mocks.NewMockConfig() store, err := client.NewStore(cfg) assert.Nil(err) assert.NotNil(store) suite.lessor, err = lease.NewLessor(store) assert.Nil(err) } func (suite *LeaseTestSuite) TearDownTest() { if suite.lessor != nil { suite.lessor.Close() } } func TestLeaseTestSuite(t *testing.T) { suite.Run(t, new(LeaseTestSuite)) }