From 3aad1b7a91a47051352d3ff121ed56a272598506 Mon Sep 17 00:00:00 2001
From: Ivan Demidov <ivan.demidov00@mail.ru>
Date: Wed, 8 Jan 2025 19:50:11 +0500
Subject: [PATCH] Add CreateOrUpdateStoreAsync

---
 .../INatsKVContext.cs                         | 11 +++++
 .../NatsKVContext.cs                          | 12 +++++
 .../KeyValueContextTest.cs                    | 48 +++++++++++++++++++
 3 files changed, 71 insertions(+)

diff --git a/src/NATS.Client.KeyValueStore/INatsKVContext.cs b/src/NATS.Client.KeyValueStore/INatsKVContext.cs
index cdf371037..2ceb0b3e7 100644
--- a/src/NATS.Client.KeyValueStore/INatsKVContext.cs
+++ b/src/NATS.Client.KeyValueStore/INatsKVContext.cs
@@ -52,6 +52,17 @@ public interface INatsKVContext
     /// <exception cref="NatsJSApiException">Server responded with an error.</exception>
     ValueTask<INatsKVStore> UpdateStoreAsync(NatsKVConfig config, CancellationToken cancellationToken = default);
 
+    /// <summary>
+    /// Creates a new Key Value Store if it doesn't exist or update if the store already exists.
+    /// </summary>
+    /// <param name="config">Key Value Store configuration</param>
+    /// <param name="cancellationToken">A <see cref="CancellationToken"/> used to cancel the API call.</param>
+    /// <returns>Key Value Store</returns>
+    /// <exception cref="NatsKVException">There was an issue with configuration</exception>
+    /// <exception cref="NatsJSException">There was an issue retrieving the response.</exception>
+    /// <exception cref="NatsJSApiException">Server responded with an error.</exception>
+    ValueTask<INatsKVStore> CreateOrUpdateStoreAsync(NatsKVConfig config, CancellationToken cancellationToken = default);
+
     /// <summary>
     /// Delete a Key Value Store
     /// </summary>
diff --git a/src/NATS.Client.KeyValueStore/NatsKVContext.cs b/src/NATS.Client.KeyValueStore/NatsKVContext.cs
index 9cbbdbc6b..646f03b01 100644
--- a/src/NATS.Client.KeyValueStore/NatsKVContext.cs
+++ b/src/NATS.Client.KeyValueStore/NatsKVContext.cs
@@ -76,6 +76,18 @@ public async ValueTask<INatsKVStore> UpdateStoreAsync(NatsKVConfig config, Cance
         return new NatsKVStore(config.Bucket, JetStreamContext, stream);
     }
 
+    /// <inheritdoc />
+    public async ValueTask<INatsKVStore> CreateOrUpdateStoreAsync(NatsKVConfig config, CancellationToken cancellationToken = default)
+    {
+        ValidateBucketName(config.Bucket);
+
+        var streamConfig = NatsKVContext.CreateStreamConfig(config);
+
+        var stream = await JetStreamContext.CreateOrUpdateStreamAsync(streamConfig, cancellationToken);
+
+        return new NatsKVStore(config.Bucket, JetStreamContext, stream);
+    }
+
     /// <inheritdoc />
     public ValueTask<bool> DeleteStoreAsync(string bucket, CancellationToken cancellationToken = default)
     {
diff --git a/tests/NATS.Client.KeyValueStore.Tests/KeyValueContextTest.cs b/tests/NATS.Client.KeyValueStore.Tests/KeyValueContextTest.cs
index 65d941ac1..438229d34 100644
--- a/tests/NATS.Client.KeyValueStore.Tests/KeyValueContextTest.cs
+++ b/tests/NATS.Client.KeyValueStore.Tests/KeyValueContextTest.cs
@@ -46,6 +46,54 @@ public async Task Update_store_test()
         updatedStatus.Info.Config.Description.Should().Be(natsKVConfig.Description);
     }
 
+    [Fact]
+    public async Task Create_store_via_create_or_update_store_test()
+    {
+        const string expectedBuketName = "kv1";
+        const string expectedDescription = "description";
+
+        var cts = new CancellationTokenSource(TimeSpan.FromSeconds(60));
+        var cancellationToken = cts.Token;
+
+        await using var server = NatsServer.StartJS();
+        await using var nats = server.CreateClientConnection();
+
+        var js = new NatsJSContext(nats);
+        var kv = new NatsKVContext(js);
+
+        var natsKVConfig = new NatsKVConfig(expectedBuketName) { Description = expectedDescription };
+        var store = await kv.CreateOrUpdateStoreAsync(natsKVConfig, cancellationToken);
+        var status = await store.GetStatusAsync(cancellationToken);
+
+        status.Bucket.Should().Be(expectedBuketName);
+        status.Info.Config.Description.Should().Be(expectedDescription);
+    }
+
+    [Fact]
+    public async Task Update_store_via_create_or_update_store_test()
+    {
+        var buketName = "kv1";
+        var expectedDescription = "Updated description";
+
+        var cts = new CancellationTokenSource(TimeSpan.FromSeconds(60));
+        var cancellationToken = cts.Token;
+
+        await using var server = NatsServer.StartJS();
+        await using var nats = server.CreateClientConnection();
+
+        var js = new NatsJSContext(nats);
+        var kv = new NatsKVContext(js);
+
+        var store = await kv.CreateStoreAsync(buketName, cancellationToken);
+        var status = await store.GetStatusAsync(cancellationToken);
+        status.Info.Config.Description.Should().BeNull();
+
+        var natsKVConfig = new NatsKVConfig(buketName) { Description = expectedDescription };
+        await kv.CreateOrUpdateStoreAsync(natsKVConfig, cancellationToken);
+        var updatedStatus = await store.GetStatusAsync(cancellationToken);
+        updatedStatus.Info.Config.Description.Should().Be(natsKVConfig.Description);
+    }
+
     [Fact]
     public async Task Delete_store_test()
     {