From 32bb7ab68e124af71d8cd16ce2d9037452f45b3d Mon Sep 17 00:00:00 2001
From: Tan Han Chong <hanchong.tan@nintex.com>
Date: Tue, 1 Oct 2024 15:28:59 +0800
Subject: [PATCH] check connection string in StorageAccountDetails since its no
 longer assigned any value in StorageConnectionString

---
 .../AzureStorageDurabilityProviderFactory.cs  |  2 +-
 .../CustomTestStorageAccountProvider.cs       | 25 +++++++++++
 ...reStorageDurabilityProviderFactoryTests.cs | 41 +++++++++++++++++++
 3 files changed, 67 insertions(+), 1 deletion(-)
 create mode 100644 test/Common/CustomTestStorageAccountProvider.cs

diff --git a/src/WebJobs.Extensions.DurableTask/AzureStorageDurabilityProviderFactory.cs b/src/WebJobs.Extensions.DurableTask/AzureStorageDurabilityProviderFactory.cs
index 242bef777..51015f79e 100644
--- a/src/WebJobs.Extensions.DurableTask/AzureStorageDurabilityProviderFactory.cs
+++ b/src/WebJobs.Extensions.DurableTask/AzureStorageDurabilityProviderFactory.cs
@@ -159,7 +159,7 @@ private AzureStorageDurabilityProvider GetAzureStorageStorageProvider(DurableCli
             // Need to check this.defaultStorageProvider != null for external clients that call GetDurabilityProvider(attribute)
             // which never initializes the defaultStorageProvider.
             if (string.Equals(this.defaultSettings?.TaskHubName, settings.TaskHubName, StringComparison.OrdinalIgnoreCase) &&
-                string.Equals(this.defaultSettings?.StorageConnectionString, settings.StorageConnectionString, StringComparison.OrdinalIgnoreCase) &&
+                string.Equals(this.defaultSettings?.StorageAccountDetails?.ConnectionString, settings.StorageAccountDetails?.ConnectionString, StringComparison.OrdinalIgnoreCase) &&
                 this.defaultStorageProvider != null)
             {
                 // It's important that clients use the same AzureStorageOrchestrationService instance
diff --git a/test/Common/CustomTestStorageAccountProvider.cs b/test/Common/CustomTestStorageAccountProvider.cs
new file mode 100644
index 000000000..3f6d3e45d
--- /dev/null
+++ b/test/Common/CustomTestStorageAccountProvider.cs
@@ -0,0 +1,25 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See LICENSE in the project root for license information.
+
+using DurableTask.AzureStorage;
+using Microsoft.WindowsAzure.Storage;
+
+namespace Microsoft.Azure.WebJobs.Extensions.DurableTask.Tests
+{
+    internal class CustomTestStorageAccountProvider : IStorageAccountProvider
+    {
+        private const string CustomConnectionString = "DefaultEndpointsProtocol=https;AccountName=test;AccountKey=dGVzdA==;EndpointSuffix=core.windows.net";
+        private readonly string customConnectionName;
+
+        public CustomTestStorageAccountProvider(string connectionName)
+        {
+            this.customConnectionName = connectionName;
+        }
+
+        public CloudStorageAccount GetCloudStorageAccount(string name) =>
+            CloudStorageAccount.Parse(name != this.customConnectionName ? TestHelpers.GetStorageConnectionString() : CustomConnectionString);
+
+        public StorageAccountDetails GetStorageAccountDetails(string name) =>
+            new StorageAccountDetails { ConnectionString = name != this.customConnectionName ? TestHelpers.GetStorageConnectionString() : CustomConnectionString };
+    }
+}
diff --git a/test/FunctionsV2/AzureStorageDurabilityProviderFactoryTests.cs b/test/FunctionsV2/AzureStorageDurabilityProviderFactoryTests.cs
index 30363c325..a55933384 100644
--- a/test/FunctionsV2/AzureStorageDurabilityProviderFactoryTests.cs
+++ b/test/FunctionsV2/AzureStorageDurabilityProviderFactoryTests.cs
@@ -154,5 +154,46 @@ public void EnvironmentIsVMSS_WorkerIdFromEnvironmentVariables()
 
             Assert.Equal("waws-prod-euapbn1-003:dw0SmallDedicatedWebWorkerRole_hr0HostRole-3-VM-13", settings.WorkerId);
         }
+
+        [Fact]
+        [Trait("Category", PlatformSpecificHelpers.TestCategory)]
+        public void CustomConnectionNameIsResolved()
+        {
+            var storageAccountProvider = new CustomTestStorageAccountProvider("CustomConnection");
+            var mockOptions = new OptionsWrapper<DurableTaskOptions>(new DurableTaskOptions());
+            var nameResolver = new Mock<INameResolver>().Object;
+
+            var factory = new AzureStorageDurabilityProviderFactory(
+                mockOptions,
+                storageAccountProvider,
+                nameResolver,
+                NullLoggerFactory.Instance,
+                TestHelpers.GetMockPlatformInformationService());
+
+            factory.GetDurabilityProvider(); // This will initialize the default connection string
+            var provider = factory.GetDurabilityProvider(new DurableClientAttribute() { ConnectionName = "CustomConnection", TaskHub = "TestHubName" });
+
+            Assert.Equal("CustomConnection", provider.ConnectionName);
+        }
+
+        [Fact]
+        [Trait("Category", PlatformSpecificHelpers.TestCategory)]
+        public void DefaultConnectionNameIsResolved()
+        {
+            var storageAccountProvider = new CustomTestStorageAccountProvider("CustomConnection");
+            var mockOptions = new OptionsWrapper<DurableTaskOptions>(new DurableTaskOptions());
+            var nameResolver = new Mock<INameResolver>().Object;
+
+            var factory = new AzureStorageDurabilityProviderFactory(
+                mockOptions,
+                storageAccountProvider,
+                nameResolver,
+                NullLoggerFactory.Instance,
+                TestHelpers.GetMockPlatformInformationService());
+
+            var provider = factory.GetDurabilityProvider();
+
+            Assert.Equal("Storage", provider.ConnectionName);
+        }
     }
 }