From bedefdac1adb90a84081c81d17d6aade738f209c Mon Sep 17 00:00:00 2001
From: Sean Kane <68240067+seankane-msft@users.noreply.github.com>
Date: Fri, 8 Jan 2021 12:40:29 -0500
Subject: [PATCH] added an option to use preparers as a back up plan for
 powershell preparer (#15945)

---
 .../devtools_testutils/powershell_preparer.py | 42 +++++++++++++------
 1 file changed, 29 insertions(+), 13 deletions(-)

diff --git a/tools/azure-sdk-tools/devtools_testutils/powershell_preparer.py b/tools/azure-sdk-tools/devtools_testutils/powershell_preparer.py
index abfbed4f8197..fd5b650885e1 100644
--- a/tools/azure-sdk-tools/devtools_testutils/powershell_preparer.py
+++ b/tools/azure-sdk-tools/devtools_testutils/powershell_preparer.py
@@ -18,6 +18,7 @@ def __init__(
         client_kwargs=None,
         random_name_enabled=False,
         use_cache=True,
+        preparers=None,
         **kwargs
     ):
         super(PowerShellPreparer, self).__init__(
@@ -29,6 +30,7 @@ def __init__(
         self.fake_values = {}
         self.real_values = {}
         self._set_secrets(**kwargs)
+        self._backup_preparers=preparers
 
     def _set_secrets(self, **kwargs):
         keys = kwargs.keys()
@@ -52,20 +54,34 @@ def create_resource(self, name, **kwargs):
 
         if self.is_live:
             self._set_mgmt_settings_real_values()
-            for key in self.needed_keys:
+            try:
+                for key in self.needed_keys:
+
+                    scrubbed_value = self.fake_values[key]
+                    if scrubbed_value:
+                        self.real_values[key.lower()] = os.environ[key.upper()]
+                        self.test_class_instance.scrubber.register_name_pair(
+                            self.real_values[key.lower()],
+                            scrubbed_value
+                        )
+                    else:
+                        template = 'To pass a live ID you must provide the scrubbed value for recordings to \
+                            prevent secrets from being written to files. {} was not given. For example: \
+                                @PowerShellPreparer("schemaregistry", schemaregistry_endpoint="fake_endpoint.servicebus.windows.net")'
+                        raise AzureTestError(template.format(key))
+            except KeyError as key_error:
+                if not self._backup_preparers:
+                    raise
+
+                self.real_values = {}
+                create_kwargs = {}
+                for preparer in self._backup_preparers:
+                    resource_name, vals = preparer._prepare_create_resource(self.test_class_instance, **create_kwargs)
+                    # vals = preparer.create_resource(name, **create_kwargs)
+                    self.real_values.update(vals)
+                    if 'resource_group' in self.real_values.keys():
+                        create_kwargs['resource_group'] = self.real_values['resource_group']
 
-                scrubbed_value = self.fake_values[key]
-                if scrubbed_value:
-                    self.real_values[key.lower()] = os.environ[key.upper()]
-                    self.test_class_instance.scrubber.register_name_pair(
-                        self.real_values[key.lower()],
-                        scrubbed_value
-                    )
-                else:
-                    template = 'To pass a live ID you must provide the scrubbed value for recordings to \
-                        prevent secrets from being written to files. {} was not given. For example: \
-                            @PowerShellPreparer("schemaregistry", schemaregistry_endpoint="fake_endpoint.servicebus.windows.net")'
-                    raise AzureTestError(template.format(key))
             return self.real_values
 
         else: