checking for azure environment variable when creating a new hub

hub.go

@@ -31,6 +31,7 @@ import (
 	"net/http"
 	"os"
 	"path"
+	"strings"
 	"sync"
 
 	"github.com/Azure/azure-amqp-common-go/v3/aad"
@@ -331,8 +332,22 @@ func hubEntryToEntity(entry *hubEntry) *HubEntity {
 }
 
 // NewHub creates a new Event Hub client for sending and receiving messages
+// NOTE: If the AZURE_ENVIRONMENT variable is set, it will be used to set the ServiceBusEndpointSuffix
+// from the corresponding azure.Environment type at the end of the namespace host string. The default
+// value is azure.PublicCloud. Azure Stack environment is currently not supported.
 func NewHub(namespace, name string, tokenProvider auth.TokenProvider, opts ...HubOption) (*Hub, error) {
-	ns, err := newNamespace(namespaceWithAzureEnvironment(namespace, tokenProvider, azure.PublicCloud))
+	var env azure.Environment
+	switch os.Getenv("AZURE_ENVIRONMENT") {
+	case strings.ToUpper(azure.ChinaCloud.Name):
+		env = azure.ChinaCloud
+	case strings.ToUpper(azure.GermanCloud.Name):
+		env = azure.GermanCloud
+	case strings.ToUpper(azure.USGovernmentCloud.Name):
+		env = azure.USGovernmentCloud
+	default:
+		env = azure.PublicCloud
+	}
+	ns, err := newNamespace(namespaceWithAzureEnvironment(namespace, tokenProvider, env))
 	if err != nil {
 		return nil, err
 	}

hub_test.go

@@ -41,6 +41,7 @@ import (
 	"github.com/Azure/azure-amqp-common-go/v3/auth"
 	"github.com/Azure/azure-amqp-common-go/v3/sas"
 	"github.com/Azure/azure-amqp-common-go/v3/uuid"
+	"github.com/Azure/go-autorest/autorest/azure"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
@@ -678,3 +679,38 @@ func (suite *eventHubSuite) captureEnv() func() {
 		suite.NoError(restoreEnv(capture))
 	}
 }
+
+func TestNewHub_withAzureEnvironmentVariable(t *testing.T) {
+	_ = os.Setenv("AZURE_ENVIRONMENT", "AZURECHINACLOUD")
+	h, err := NewHub("test", "test", &aad.TokenProvider{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !strings.HasSuffix(h.namespace.host, azure.ChinaCloud.ServiceBusEndpointSuffix) {
+		t.Fatalf("did not set appropriate endpoint suffix. Expected: %v, Received: %v", azure.ChinaCloud.ServiceBusEndpointSuffix, h.namespace.host)
+	}
+	_ = os.Setenv("AZURE_ENVIRONMENT", "AZUREGERMANCLOUD")
+	h, err = NewHub("test", "test", &aad.TokenProvider{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !strings.HasSuffix(h.namespace.host, azure.GermanCloud.ServiceBusEndpointSuffix) {
+		t.Fatalf("did not set appropriate endpoint suffix. Expected: %v, Received: %v", azure.GermanCloud.ServiceBusEndpointSuffix, h.namespace.host)
+	}
+	_ = os.Setenv("AZURE_ENVIRONMENT", "AZUREUSGOVERNMENTCLOUD")
+	h, err = NewHub("test", "test", &aad.TokenProvider{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !strings.HasSuffix(h.namespace.host, azure.USGovernmentCloud.ServiceBusEndpointSuffix) {
+		t.Fatalf("did not set appropriate endpoint suffix. Expected: %v, Received: %v", azure.USGovernmentCloud.ServiceBusEndpointSuffix, h.namespace.host)
+	}
+	_ = os.Unsetenv("AZURE_ENVIRONMENT")
+	h, err = NewHub("test", "test", &aad.TokenProvider{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !strings.HasSuffix(h.namespace.host, azure.PublicCloud.ServiceBusEndpointSuffix) {
+		t.Fatalf("did not set appropriate endpoint suffix. Expected: %v, Received: %v", azure.PublicCloud.ServiceBusEndpointSuffix, h.namespace.host)
+	}
+}