forge/modules/globallock/locker_test.go

// Copyright 2026 Hanzo AI, Inc. All rights reserved.
// SPDX-License-Identifier: MIT

package globallock

import (
	"context"
	"os"
	"sync"
	"testing"
	"time"

	"github.com/hanzoai/git/modules/test"
	"github.com/hanzoai/git/modules/util"

	"github.com/hanzokv/lock/v4"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func newTestKVLocker(t *testing.T) Locker {
	t.Helper()
	kvURL := util.IfZero(os.Getenv("TEST_KV_URL"), "kv://127.0.0.1:6379/0")
	rl := NewKVLocker(kvURL).(*kvLocker)
	err := rl.conn.Ping(t.Context()).Err()
	if err != nil && test.AllowSkipExternalService() {
		t.Skip("no KV server for testing, skipped")
	}
	require.NoError(t, err, "KV error for testing: %v", err)
	return rl
}

func TestLocker(t *testing.T) {
	t.Run("kv", func(t *testing.T) {
		defer test.MockVariableValue(&kvLockExpiry, 5*time.Second)() // make it shorter for testing
		locker := newTestKVLocker(t)
		testLocker(t, locker)
		testKVLocker(t, locker.(*kvLocker))
		require.NoError(t, locker.(*kvLocker).Close())
	})
	t.Run("memory", func(t *testing.T) {
		locker := NewMemoryLocker()
		testLocker(t, locker)
		testMemoryLocker(t, locker.(*memoryLocker))
	})
}

func testLocker(t *testing.T, locker Locker) {
	t.Run("lock", func(t *testing.T) {
		parentCtx := t.Context()
		release, err := locker.Lock(parentCtx, "test")
		defer release()

		assert.NoError(t, err)

		func() {
			ctx, cancel := context.WithTimeout(t.Context(), time.Second)
			defer cancel()
			release, err := locker.Lock(ctx, "test")
			defer release()

			assert.Error(t, err)
		}()

		release()

		func() {
			release, err := locker.Lock(t.Context(), "test")
			defer release()

			assert.NoError(t, err)
		}()
	})

	t.Run("try lock", func(t *testing.T) {
		parentCtx := t.Context()
		ok, release, err := locker.TryLock(parentCtx, "test")
		defer release()

		assert.True(t, ok)
		assert.NoError(t, err)

		func() {
			ctx, cancel := context.WithTimeout(t.Context(), time.Second)
			defer cancel()
			ok, release, err := locker.TryLock(ctx, "test")
			defer release()

			assert.False(t, ok)
			assert.NoError(t, err)
		}()

		release()

		func() {
			ok, release, _ := locker.TryLock(t.Context(), "test")
			defer release()

			assert.True(t, ok)
		}()
	})

	t.Run("wait and acquired", func(t *testing.T) {
		ctx := t.Context()
		release, err := locker.Lock(ctx, "test")
		require.NoError(t, err)

		wg := &sync.WaitGroup{}
		wg.Go(func() {
			started := time.Now()
			release, err := locker.Lock(t.Context(), "test") // should be blocked for seconds
			defer release()
			assert.Greater(t, time.Since(started), time.Second)
			assert.NoError(t, err)
		})

		time.Sleep(2 * time.Second)
		release()

		wg.Wait()
	})

	t.Run("multiple release", func(t *testing.T) {
		ctx := t.Context()

		release1, err := locker.Lock(ctx, "test")
		require.NoError(t, err)

		release1()

		release2, err := locker.Lock(ctx, "test")
		defer release2()
		require.NoError(t, err)

		// Call release1 again,
		// it should not panic or block,
		// and it shouldn't affect the other lock
		release1()

		ok, release3, err := locker.TryLock(ctx, "test")
		defer release3()
		require.NoError(t, err)
		// It should be able to acquire the lock;
		// otherwise, it means the lock has been released by release1
		assert.False(t, ok)
	})
}

// testMemoryLocker does specific tests for memoryLocker
func testMemoryLocker(t *testing.T, locker *memoryLocker) {
	// nothing to do
}

// testKVLocker does specific tests for kvLocker
func testKVLocker(t *testing.T, locker *kvLocker) {
	defer func() {
		// This case should be tested at the end.
		// Otherwise, it will affect other tests.
		t.Run("close", func(t *testing.T) {
			assert.NoError(t, locker.Close())
			_, err := locker.Lock(t.Context(), "test")
			assert.Error(t, err)
		})
	}()

	t.Run("failed extend", func(t *testing.T) {
		release, err := locker.Lock(t.Context(), "test")
		defer release()
		require.NoError(t, err)

		// It simulates that there are some problems with extending like network issues or KV server down.
		v, ok := locker.mutexM.Load("test")
		require.True(t, ok)
		m := v.(*redsync.Mutex)
		_, _ = m.Unlock() // release it to make it impossible to extend

		// In current design, callers can't know the lock can't be extended.
		// Just keep this case to improve the test coverage.
	})
}