From 8f111aca694e4640b448bbb882f0612f4a43377b Mon Sep 17 00:00:00 2001
From: sakno <roman.sakno@gmail.com>
Date: Sun, 16 Jun 2024 21:00:46 +0300
Subject: [PATCH] Added tests for LeaseConsumer

---
 .../Threading/Leases/LeaseTests.cs            | 114 +++++++++++++++++-
 1 file changed, 111 insertions(+), 3 deletions(-)

diff --git a/src/DotNext.Tests/Threading/Leases/LeaseTests.cs b/src/DotNext.Tests/Threading/Leases/LeaseTests.cs
index b73ee9e71..7e4475618 100644
--- a/src/DotNext.Tests/Threading/Leases/LeaseTests.cs
+++ b/src/DotNext.Tests/Threading/Leases/LeaseTests.cs
@@ -7,6 +7,7 @@ public sealed class LeaseTests : Test
     public static async Task AcquireOrRenewInitialState()
     {
         using var provider = new TestLeaseProvider(DefaultTimeout);
+        Equal(DefaultTimeout, provider.TimeToLive);
         Null(await provider.TryRenewAsync(default, true));
         Null(await provider.TryRenewAsync(default, false));
         Null(await provider.ReleaseAsync(default));
@@ -16,14 +17,37 @@ public static async Task AcquireOrRenewInitialState()
     }
 
     [Fact]
-    public static async Task Acquisition()
+    public static async Task AcquireRelease()
     {
         using var provider = new TestLeaseProvider(DefaultTimeout);
-        NotNull(await provider.TryAcquireAsync());
+        var result = NotNull(await provider.TryAcquireAsync());
+        NotNull(await provider.ReleaseAsync(result.State.Identity));
+
+        await provider.AcquireAsync();
+        Null(await provider.TryAcquireAsync());
+        NotNull(await provider.UnsafeTryReleaseAsync());
+    }
+
+    [Fact]
+    public static async Task RenewOrAcquire()
+    {
+        using var provider = new TestLeaseProvider(DefaultTimeout);
+        var result = NotNull(await provider.TryAcquireOrRenewAsync(default));
+        var result2 = NotNull(await provider.TryAcquireOrRenewAsync(result.State.Identity));
+        True(result.State.Identity << result2.State.Identity);
     }
 
     [Fact]
-    public static void Preceding()
+    public static async Task RenewAfterRelease()
+    {
+        using var provider = new TestLeaseProvider(DefaultTimeout);
+        var result = await provider.AcquireAsync();
+        await provider.UnsafeReleaseAsync();
+        Null(await provider.TryRenewAsync(result.State.Identity, reacquire: false));
+    }
+
+    [Fact]
+    public static void Precedence()
     {
         True(default(LeaseIdentity) << new LeaseIdentity { Version = 1UL });
         False(default(LeaseIdentity) >> new LeaseIdentity { Version = 1UL });
@@ -43,6 +67,62 @@ public static async Task FightForLease()
         True(tasks is [null, not null] or [not null, null]);
     }
 
+    [Fact]
+    public static async Task FightForLeaseUsingConsumer()
+    {
+        using var provider = new TestLeaseProvider(DefaultTimeout);
+        await using var consumer1 = new TestLeaseConsumer(provider);
+        await using var consumer2 = new TestLeaseConsumer(provider);
+
+        var acquisition1 = Task.Run(async () => await consumer1.TryAcquireAsync());
+        var acquisition2 = Task.Run(async () => await consumer1.TryAcquireAsync());
+
+        var tasks = await Task.WhenAll(acquisition1, acquisition2);
+
+        True(tasks is [{ IsExpired: false }, { IsExpired: true }] or [{ IsExpired: true }, { IsExpired: false }]);
+    }
+
+    [Fact]
+    public static async Task ConsumerTokenState()
+    {
+        using var provider = new TestLeaseProvider(TimeSpan.FromMilliseconds(100));
+        await using var consumer = new TestLeaseConsumer(provider);
+        True(consumer.Token.IsCancellationRequested);
+
+        False((await consumer.TryAcquireAsync()).IsExpired);
+        False(consumer.Token.IsCancellationRequested);
+
+        await Task.Delay(150);
+        True(consumer.Token.IsCancellationRequested);
+        False(await consumer.ReleaseAsync());
+    }
+
+    [Fact]
+    public static async Task ConsumerRenew()
+    {
+        using var provider = new TestLeaseProvider(DefaultTimeout);
+        await using var consumer = new TestLeaseConsumer(provider);
+        await consumer.TryAcquireAsync();
+        var expected = consumer.Token;
+
+        await consumer.TryRenewAsync();
+        Equal(consumer.Token, expected);
+    }
+
+    [Fact]
+    public static async Task DisposeConcurrently()
+    {
+        Task task;
+        using (var provider = new TestLeaseProvider(DefaultTimeout))
+        {
+            await provider.AcquireAsync();
+            task = provider.AcquireAsync().AsTask();
+            False(task.IsCompleted);
+        }
+
+        await ThrowsAsync<ObjectDisposedException>(Func.Constant(task));
+    }
+
     private sealed class TestLeaseProvider(TimeSpan ttl) : LeaseProvider<int>(ttl)
     {
         private readonly AsyncReaderWriterLock syncRoot = new();
@@ -90,4 +170,32 @@ protected override void Dispose(bool disposing)
             base.Dispose(disposing);
         }
     }
+    
+    private sealed class TestLeaseConsumer(TestLeaseProvider provider) : LeaseConsumer
+    {
+        protected override async ValueTask<AcquisitionResult?> TryAcquireCoreAsync(CancellationToken token = default)
+        {
+            return await provider.TryAcquireAsync(token) is { } result
+                ? new AcquisitionResult { Identity = result.State.Identity, TimeToLive = result.Lifetime.Value }
+                : null;
+        }
+
+        protected override async ValueTask<AcquisitionResult?> TryRenewCoreAsync(LeaseIdentity identity,
+            CancellationToken token)
+        {
+            return await provider.TryAcquireOrRenewAsync(identity, token) is { } result
+                ? new AcquisitionResult { Identity = result.State.Identity, TimeToLive = result.Lifetime.Value }
+                : null;
+        }
+
+        protected override ValueTask<LeaseIdentity?> ReleaseCoreAsync(LeaseIdentity identity,
+            CancellationToken token)
+            => provider.ReleaseAsync(identity, token);
+
+        protected override ValueTask DisposeAsyncCore()
+        {
+            provider.Dispose();
+            return base.DisposeAsyncCore();
+        }
+    }
 }
\ No newline at end of file