diff --git a/hub.go b/hub.go index bd075de..515e3fc 100644 --- a/hub.go +++ b/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 } diff --git a/hub_test.go b/hub_test.go index c38e72f..11cb42e 100644 --- a/hub_test.go +++ b/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) + } +}