batch_image_download_limiter_test.go 1.1 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243
  1. //go:build unit
  2. package repository
  3. import (
  4. "context"
  5. "testing"
  6. "time"
  7. "github.com/Wei-Shaw/sub2api/internal/service"
  8. "github.com/alicebob/miniredis/v2"
  9. "github.com/redis/go-redis/v9"
  10. "github.com/stretchr/testify/require"
  11. )
  12. func TestBatchImageDownloadLimiter_AcquireDenyReleaseAndTTL(t *testing.T) {
  13. ctx := context.Background()
  14. mr := miniredis.RunT(t)
  15. rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
  16. t.Cleanup(func() { _ = rdb.Close() })
  17. limiter := &batchImageDownloadLimiter{
  18. rdb: rdb,
  19. activePrefix: defaultBatchImageDownloadActivePrefix,
  20. maxActive: 1,
  21. ttl: time.Minute,
  22. }
  23. permit, err := limiter.Acquire(ctx, "11", "zip")
  24. require.NoError(t, err)
  25. require.NotNil(t, permit)
  26. require.True(t, mr.TTL(limiter.activeKey("11")) > 0)
  27. _, err = limiter.Acquire(ctx, "11", "zip")
  28. require.ErrorIs(t, err, service.ErrBatchImageDownloadLimited)
  29. require.NoError(t, permit.Release(ctx))
  30. require.NoError(t, permit.Release(ctx))
  31. require.False(t, mr.Exists(limiter.activeKey("11")))
  32. permit, err = limiter.Acquire(ctx, "11", "zip")
  33. require.NoError(t, err)
  34. require.NotNil(t, permit)
  35. }