package distlock import ( "context" "fmt" "sync" "sync/atomic" "testing" "time" "github.com/stretchr/testify/require" ) func TestMemoryStoreAcquire(t *testing.T) { store := newTestMemoryStore() ctx := context.Background() acquired, err := store.Acquire(ctx, "resource-a", []byte("owner-a"), time.Minute) require.NoError(t, err) require.True(t, acquired) acquired, err = store.Acquire(ctx, "resource-a", []byte("owner-b"), time.Minute) require.NoError(t, err) require.False(t, acquired) acquired, err = store.Acquire(ctx, "resource-b", []byte("owner-b"), time.Minute) require.NoError(t, err) require.True(t, acquired) } func TestMemoryStoreExpirationAndStaleOwner(t *testing.T) { now := time.Date(2026, time.August, 26, 12, 0, 0, 0, time.UTC) store := newTestMemoryStore() store.now = func() time.Time { return now } ctx := context.Background() oldToken := []byte("old-owner") newToken := []byte("new-owner") acquired, err := store.Acquire(ctx, "resource", oldToken, time.Minute) require.NoError(t, err) require.True(t, acquired) // A lease is expired at its expiration time, not only after it. now = now.Add(time.Minute) acquired, err = store.Acquire(ctx, "resource", newToken, time.Minute) require.NoError(t, err) require.True(t, acquired) released, err := store.Release(ctx, "resource", oldToken) require.NoError(t, err) require.False(t, released) renewed, err := store.Renew(ctx, "resource", oldToken, time.Minute) require.NoError(t, err) require.False(t, renewed) released, err = store.Release(ctx, "resource", newToken) require.NoError(t, err) require.True(t, released) } func TestMemoryStoreRenew(t *testing.T) { now := time.Date(2026, time.August, 26, 12, 0, 0, 0, time.UTC) store := newTestMemoryStore() store.now = func() time.Time { return now } ctx := context.Background() token := []byte("owner") acquired, err := store.Acquire(ctx, "resource", token, time.Minute) require.NoError(t, err) require.True(t, acquired) now = now.Add(30 * time.Second) renewed, err := store.Renew(ctx, "resource", []byte("another-owner"), time.Minute) require.NoError(t, err) require.False(t, renewed) renewed, err = store.Renew(ctx, "resource", token, time.Minute) require.NoError(t, err) require.True(t, renewed) // The renewal extends the lease from the renewal time. now = now.Add(30 * time.Second) acquired, err = store.Acquire(ctx, "resource", []byte("another-owner"), time.Minute) require.NoError(t, err) require.False(t, acquired) now = now.Add(30 * time.Second) renewed, err = store.Renew(ctx, "resource", token, time.Minute) require.NoError(t, err) require.False(t, renewed) } func TestMemoryStoreRelease(t *testing.T) { store := newTestMemoryStore() ctx := context.Background() token := []byte("owner") released, err := store.Release(ctx, "resource", token) require.NoError(t, err) require.False(t, released) acquired, err := store.Acquire(ctx, "resource", token, time.Minute) require.NoError(t, err) require.True(t, acquired) released, err = store.Release(ctx, "resource", []byte("another-owner")) require.NoError(t, err) require.False(t, released) released, err = store.Release(ctx, "resource", token) require.NoError(t, err) require.True(t, released) released, err = store.Release(ctx, "resource", token) require.NoError(t, err) require.False(t, released) } func TestMemoryStoreCopiesToken(t *testing.T) { store := newTestMemoryStore() ctx := context.Background() token := []byte("owner") originalToken := append([]byte(nil), token...) acquired, err := store.Acquire(ctx, "resource", token, time.Minute) require.NoError(t, err) require.True(t, acquired) token[0] = 'x' released, err := store.Release(ctx, "resource", originalToken) require.NoError(t, err) require.True(t, released) } func TestMemoryStoreValidation(t *testing.T) { tests := []struct { name string run func(*MemoryStore) error }{ { name: "acquire with empty resource", run: func(store *MemoryStore) error { _, err := store.Acquire(context.Background(), "", []byte("owner"), time.Minute) return err }, }, { name: "acquire with empty token", run: func(store *MemoryStore) error { _, err := store.Acquire(context.Background(), "resource", nil, time.Minute) return err }, }, { name: "acquire with zero TTL", run: func(store *MemoryStore) error { _, err := store.Acquire(context.Background(), "resource", []byte("owner"), 0) return err }, }, { name: "renew with negative TTL", run: func(store *MemoryStore) error { _, err := store.Renew(context.Background(), "resource", []byte("owner"), -time.Second) return err }, }, { name: "release with empty token", run: func(store *MemoryStore) error { _, err := store.Release(context.Background(), "resource", nil) return err }, }, { name: "cancelled context", run: func(store *MemoryStore) error { ctx, cancel := context.WithCancel(context.Background()) cancel() _, err := store.Acquire(ctx, "resource", []byte("owner"), time.Minute) return err }, }, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { store := newTestMemoryStore() require.Error(t, test.run(store)) // Invalid operations must not create or replace a lease. acquired, err := store.Acquire(context.Background(), "resource", []byte("valid-owner"), time.Minute) require.NoError(t, err) require.True(t, acquired) }) } } func TestMemoryStoreConcurrentAcquire(t *testing.T) { const attempts = 100 store := newTestMemoryStore() ctx := context.Background() start := make(chan struct{}) errCh := make(chan error, attempts) var acquired atomic.Int64 var wg sync.WaitGroup for i := range attempts { wg.Add(1) go func() { defer wg.Done() <-start ok, err := store.Acquire(ctx, "resource", []byte(fmt.Sprintf("owner-%d", i)), time.Minute) if err != nil { errCh <- err return } if ok { acquired.Add(1) } }() } close(start) wg.Wait() close(errCh) for err := range errCh { require.NoError(t, err) } require.EqualValues(t, 1, acquired.Load()) } func newTestMemoryStore() *MemoryStore { store := NewMemoryStore() store.now = func() time.Time { return time.Date(2026, time.August, 26, 12, 0, 0, 0, time.UTC) } return store }