Conditionally allow and test connection strings (only in TME) (#24047)

Local resources can no longer contain any connection strings so we need to enable it conditionally. A future change in New-TestResources.ps1 will pass the flag, so this is just making the change ahead of time.
Azure · Feb 4, 2025 · d6bd6b7 · d6bd6b7
1 parent a502d23
commit d6bd6b7
Show file tree

Hide file tree

Showing 7 changed files with 82 additions and 46 deletions.
diff --git a/sdk/messaging/azeventhubs/internal/test/test_helpers.go b/sdk/messaging/azeventhubs/internal/test/test_helpers.go
@@ -106,6 +106,12 @@ type ConnectionParamsForTest struct {
 }
 
 func (c ConnectionParamsForTest) CS(t *testing.T) struct{ Primary, ListenOnly, SendOnly, Storage string } {
+	if val, exists := os.LookupEnv("EVENTHUB_CONNECTION_STRING_LISTEN_ONLY"); exists && val == "" {
+		// This happens if we're not in the TME subscription - the variable will just be set to an empty string
+		// rather than not existing, altogether.
+		t.Skip("Not in TME, skipping connection string tests")
+	}
+
 	envVars := mustGetEnvironmentVars(t, []string{
 		"EVENTHUB_CONNECTION_STRING_LISTEN_ONLY",
 		"EVENTHUB_CONNECTION_STRING_SEND_ONLY",

diff --git a/sdk/messaging/azeventhubs/test-resources.bicep b/sdk/messaging/azeventhubs/test-resources.bicep
@@ -11,15 +11,15 @@ param storageEndpointSuffix string = environment().suffixes.storage
 @description('The resource location')
 param location string = resourceGroup().location
 
+param tenantIsTME bool = false
+
 var apiVersion = '2017-04-01'
-var storageApiVersion = '2019-04-01'
 var namespaceName = baseName
 var storageAccountName = 'storage${baseName}'
 var containerName = 'container'
-var iotName = 'iot${baseName}'
 var authorizationName = '${baseName}/RootManageSharedAccessKey'
 
-resource namespace 'Microsoft.EventHub/namespaces@2017-04-01' = {
+resource namespace 'Microsoft.EventHub/namespaces@2024-01-01' = {
   name: namespaceName
   location: location
   sku: {
@@ -28,6 +28,7 @@ resource namespace 'Microsoft.EventHub/namespaces@2017-04-01' = {
     capacity: 5
   }
   properties: {
+    disableLocalAuth: !tenantIsTME
     isAutoInflateEnabled: false
     maximumThroughputUnits: 0
   }
@@ -146,18 +147,25 @@ output EVENTHUB_NAME string = eventHub.name
 output EVENTHUB_LINKSONLY_NAME string = linksonly.name
 
 // connection strings
-output EVENTHUB_CONNECTION_STRING string = listKeys(
-  resourceId('Microsoft.EventHub/namespaces/authorizationRules', namespaceName, 'RootManageSharedAccessKey'),
-  apiVersion
-).primaryConnectionString
-output EVENTHUB_CONNECTION_STRING_LISTEN_ONLY string = listKeys(
-  resourceId('Microsoft.EventHub/namespaces/authorizationRules', namespaceName, authorizedListenOnly.name),
-  apiVersion
-).primaryConnectionString
-output EVENTHUB_CONNECTION_STRING_SEND_ONLY string = listKeys(
-  resourceId('Microsoft.EventHub/namespaces/authorizationRules', namespaceName, authorizedSendOnly.name),
-  apiVersion
-).primaryConnectionString
+output EVENTHUB_CONNECTION_STRING string = tenantIsTME
+  ? listKeys(
+      resourceId('Microsoft.EventHub/namespaces/authorizationRules', namespaceName, 'RootManageSharedAccessKey'),
+      apiVersion
+    ).primaryConnectionString
+  : ''
+
+output EVENTHUB_CONNECTION_STRING_LISTEN_ONLY string = tenantIsTME
+  ? listKeys(
+      resourceId('Microsoft.EventHub/namespaces/authorizationRules', namespaceName, authorizedListenOnly.name),
+      apiVersion
+    ).primaryConnectionString
+  : ''
+output EVENTHUB_CONNECTION_STRING_SEND_ONLY string = tenantIsTME
+  ? listKeys(
+      resourceId('Microsoft.EventHub/namespaces/authorizationRules', namespaceName, authorizedSendOnly.name),
+      apiVersion
+    ).primaryConnectionString
+  : ''
 
 output RESOURCE_GROUP string = resourceGroup().name
 output AZURE_SUBSCRIPTION_ID string = subscription().subscriptionId
diff --git a/sdk/messaging/azservicebus/admin/admin_client_test.go b/sdk/messaging/azservicebus/admin/admin_client_test.go
@@ -24,7 +24,6 @@ import (
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
 	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/atom"
 	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/auth"
-	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/conn"
 	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/test"
 	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus/internal/utils"
 	"github.com/stretchr/testify/require"
@@ -721,10 +720,8 @@ func TestAdminClient_Forwarding(t *testing.T) {
 
 	defer deleteTopic(t, adminClient, topicName)
 
-	cs := test.MustGetEnvVar(t, test.EnvKeyConnectionString)
-	parsed, err := conn.ParseConnectionString(cs)
-	require.NoError(t, err)
-	forwardToName := to.Ptr(fmt.Sprintf("sb://%s/%s", parsed.FullyQualifiedNamespace, forwardToQueueName))
+	endpoint := test.MustGetEnvVar(t, test.EnvKeyEndpoint)
+	forwardToName := to.Ptr(fmt.Sprintf("sb://%s/%s", endpoint, forwardToQueueName))
 
 	_, err = adminClient.CreateSubscription(context.Background(), topicName, "sub1", &CreateSubscriptionOptions{
 		Properties: &SubscriptionProperties{
@@ -1985,8 +1982,8 @@ func setupLowPrivTest(t *testing.T) *struct {
 	Cleanup   func()
 } {
 	adminClient := newAdminClientForTest(t, nil)
-
-	lowPrivAdminClient, err := NewClientFromConnectionString(test.MustGetEnvVar(t, test.EnvKeyConnectionStringNoManage), nil) // allowed connection string
+	cs := test.GetConnectionString(t, test.EnvKeyConnectionStringNoManage)
+	lowPrivAdminClient, err := NewClientFromConnectionString(cs, nil) // allowed connection string
 	require.NoError(t, err)
 
 	nanoSeconds := time.Now().UnixNano()

diff --git a/sdk/messaging/azservicebus/client_test.go b/sdk/messaging/azservicebus/client_test.go
@@ -104,8 +104,8 @@ func TestNewClientWithWebsockets(t *testing.T) {
 
 func TestNewClientUsingSharedAccessSignature(t *testing.T) {
 	getLogsFn := test.CaptureLogsForTest(false)
+	cs := test.GetConnectionString(t, test.EnvKeyConnectionString)
 
-	cs := test.MustGetEnvVar(t, test.EnvKeyConnectionString)
 	sasCS, err := sas.CreateConnectionStringWithSASUsingExpiry(cs, time.Now().UTC().Add(time.Hour))
 	require.NoError(t, err)
 
@@ -328,7 +328,7 @@ func TestClientPropagatesRetryOptionsForSessions(t *testing.T) {
 }
 
 func TestClientUnauthorizedCreds(t *testing.T) {
-	allPowerfulCS := test.MustGetEnvVar(t, test.EnvKeyConnectionString)
+	allPowerfulCS := test.GetConnectionString(t, test.EnvKeyConnectionString)
 	queueName := "testqueue"
 
 	t.Run("ListenOnly with Sender", func(t *testing.T) {

diff --git a/sdk/messaging/azservicebus/internal/test/test_helpers.go b/sdk/messaging/azservicebus/internal/test/test_helpers.go
@@ -92,11 +92,13 @@ func MustGetEnvVars[KeyT ~string](keys []KeyT) map[KeyT]string {
 	var missing []string
 
 	for _, k := range keys {
-		m[k] = os.Getenv(string(k))
+		v, exists := os.LookupEnv(string(k))
 
-		if m[k] == "" {
+		if !exists {
 			missing = append(missing, string(k))
 		}
+
+		m[k] = v
 	}
 
 	if len(missing) != 0 {
@@ -107,6 +109,23 @@ func MustGetEnvVars[KeyT ~string](keys []KeyT) map[KeyT]string {
 	return m
 }
 
+func GetConnectionString(t *testing.T, name EnvKey) string {
+	if recording.GetRecordMode() == recording.PlaybackMode {
+		t.Skipf("Skipping live test when in recording.PlaybackMode")
+		return ""
+	}
+
+	val, exists := os.LookupEnv(string(name))
+
+	if exists && val == "" {
+		// This happens if we're not in the TME subscription - the variable will just be set to an empty string
+		// rather than not existing, altogether.
+		t.Skip("Not in TME, skipping connection string tests")
+	}
+
+	return val
+}
+
 func MustGetEnvVar(t *testing.T, name EnvKey) string {
 	if recording.GetRecordMode() == recording.PlaybackMode {
 		t.Skipf("Skipping live test when in recording.PlaybackMode")
@@ -141,10 +160,9 @@ func NewClient[OptionsT any, ClientT any](t *testing.T, args NewClientArgs[Optio
 	}
 
 	if options.UseConnectionString {
+		cs := GetConnectionString(t, EnvKeyConnectionString)
 		env := MustGetEnvVars([]EnvKey{EnvKeyConnectionString, EnvKeyConnectionStringPremium})
 
-		cs := env[EnvKeyConnectionString]
-
 		if options.UsePremium {
 			cs = env[EnvKeyConnectionStringPremium]
 		}
@@ -170,8 +188,9 @@ func NewClient[OptionsT any, ClientT any](t *testing.T, args NewClientArgs[Optio
 }
 
 func CreateExpiringQueue(t *testing.T, qd *atom.QueueDescription) (string, func()) {
-	cs := MustGetEnvVar(t, EnvKeyConnectionString)
-	em, err := atom.NewEntityManagerWithConnectionString(cs, "", nil)
+	iv := GetIdentityVars(t)
+
+	em, err := atom.NewEntityManager(iv.Endpoint, iv.Cred, "", nil)
 	require.NoError(t, err)
 
 	queueName := RandomString("queue", 5)

diff --git a/sdk/messaging/azservicebus/receiver_test.go b/sdk/messaging/azservicebus/receiver_test.go
@@ -823,11 +823,11 @@ func TestReceiverMessageLockExpires(t *testing.T) {
 }
 
 func TestReceiverUnauthorizedCreds(t *testing.T) {
-	allPowerfulCS := test.MustGetEnvVar(t, test.EnvKeyConnectionString)
+	allPowerfulCS := test.GetConnectionString(t, test.EnvKeyConnectionString)
 	queueName := "testqueue"
 
 	t.Run("ListenOnly with Sender", func(t *testing.T) {
-		cs := test.MustGetEnvVar(t, test.EnvKeyConnectionStringListenOnly)
+		cs := test.GetConnectionString(t, test.EnvKeyConnectionStringListenOnly)
 
 		client, err := NewClientFromConnectionString(cs, nil) // allowed connection string
 		require.NoError(t, err)

diff --git a/sdk/messaging/azservicebus/test-resources.bicep b/sdk/messaging/azservicebus/test-resources.bicep
@@ -16,30 +16,36 @@ param enablePremium bool = true
 @description('Disable applying any RBAC rules')
 param disableAddingRBACRole bool = false
 
+param tenantIsTME bool = false
+
 var apiVersion = '2017-04-01'
 var serviceBusDataOwnerRoleId = '/subscriptions/${subscription().subscriptionId}/providers/Microsoft.Authorization/roleDefinitions/090c5cfd-751d-490a-894a-3ce6f1109419'
 
 var sbPremiumName = 'sb-premium-${baseName}'
 
-resource servicebus 'Microsoft.ServiceBus/namespaces@2018-01-01-preview' = {
+resource servicebus 'Microsoft.ServiceBus/namespaces@2024-01-01' = {
   name: baseName
-  location: location
+  location: location 
   sku: {
     name: 'Standard'
     tier: 'Standard'
   }
   properties: {
+    disableLocalAuth: !tenantIsTME
     zoneRedundant: false
   }
 }
 
-resource servicebusPremium 'Microsoft.ServiceBus/namespaces@2018-01-01-preview' = if (enablePremium) {
+resource servicebusPremium 'Microsoft.ServiceBus/namespaces@2024-01-01' = if (enablePremium) {
   name: sbPremiumName
   location: location
   sku: {
     name: 'Premium'
     tier: 'Premium'
   }
+  properties: {
+    disableLocalAuth: !tenantIsTME    
+  }
 }
 
 resource authorizationRuleName 'Microsoft.ServiceBus/namespaces/AuthorizationRules@2015-08-01' = {
@@ -136,31 +142,31 @@ resource testQueueWithSessions 'Microsoft.ServiceBus/namespaces/queues@2017-04-0
   }
 }
 
-output SERVICEBUS_CONNECTION_STRING string = listKeys(
+output SERVICEBUS_CONNECTION_STRING string = tenantIsTME ? listKeys(
   resourceId('Microsoft.ServiceBus/namespaces/authorizationRules', baseName, 'RootManageSharedAccessKey'),
   apiVersion
-).primaryConnectionString
+).primaryConnectionString : ''
 
 // connection strings with fewer rights - no manage rights, listen only (ie, receive) and send only.
-output SERVICEBUS_CONNECTION_STRING_NO_MANAGE string = listKeys(
+output SERVICEBUS_CONNECTION_STRING_NO_MANAGE string = tenantIsTME ? listKeys(
   resourceId('Microsoft.ServiceBus/namespaces/authorizationRules', baseName, 'NoManage'),
   apiVersion
-).primaryConnectionString
-output SERVICEBUS_CONNECTION_STRING_SEND_ONLY string = listKeys(
+).primaryConnectionString : ''
+output SERVICEBUS_CONNECTION_STRING_SEND_ONLY string = tenantIsTME ? listKeys(
   resourceId('Microsoft.ServiceBus/namespaces/authorizationRules', baseName, 'SendOnly'),
   apiVersion
-).primaryConnectionString
-output SERVICEBUS_CONNECTION_STRING_LISTEN_ONLY string = listKeys(
+).primaryConnectionString : ''
+output SERVICEBUS_CONNECTION_STRING_LISTEN_ONLY string = tenantIsTME ? listKeys(
   resourceId('Microsoft.ServiceBus/namespaces/authorizationRules', baseName, 'ListenOnly'),
   apiVersion
-).primaryConnectionString
+).primaryConnectionString: ''
 
-output SERVICEBUS_CONNECTION_STRING_PREMIUM string = enablePremium
+output SERVICEBUS_CONNECTION_STRING_PREMIUM string = enablePremium && tenantIsTME
   ? listKeys(
       resourceId('Microsoft.ServiceBus/namespaces/authorizationRules', sbPremiumName, 'RootManageSharedAccessKey'),
       apiVersion
     ).primaryConnectionString
-  : '<premium endpoint not deployed'
+  : ''
 
 output SERVICEBUS_ENDPOINT string = replace(
   replace(servicebus.properties.serviceBusEndpoint, ':443/', ''),
@@ -169,7 +175,7 @@ output SERVICEBUS_ENDPOINT string = replace(
 )
 output SERVICEBUS_ENDPOINT_PREMIUM string = enablePremium
   ? replace(replace(servicebusPremium.properties.serviceBusEndpoint, ':443/', ''), 'https://', '')
-  : '<premium endpoint not deployed'
+  : ''
 
 output QUEUE_NAME string = 'testQueue'
 output QUEUE_NAME_WITH_SESSIONS string = 'testQueueWithSessions'