From 542796d3a7c518bb34c9ea24ee761667df3840c1 Mon Sep 17 00:00:00 2001
From: seankane-msft <seankane@microsoft.com>
Date: Tue, 29 Dec 2020 12:47:44 -0500
Subject: [PATCH] added an option to use preparers as a back up plan for
 powershell preparer

---
 .../tests/test_table_client_cosmos.py         |  7 +++-
 .../devtools_testutils/powershell_preparer.py | 42 +++++++++++++------
 2 files changed, 35 insertions(+), 14 deletions(-)

diff --git a/sdk/tables/azure-data-tables/tests/test_table_client_cosmos.py b/sdk/tables/azure-data-tables/tests/test_table_client_cosmos.py
index 56e43a7259b0..692b6ac40ae2 100644
--- a/sdk/tables/azure-data-tables/tests/test_table_client_cosmos.py
+++ b/sdk/tables/azure-data-tables/tests/test_table_client_cosmos.py
@@ -28,7 +28,12 @@
 CosmosPSPreparer = functools.partial(
     PowerShellPreparer, "tables",
     tables_cosmos_account_name="fake_cosmos_account",
-    tables_primary_cosmos_account_key="fakecosmosaccountkey")
+    tables_primary_cosmos_account_key="fakecosmosaccountkey",
+    preparers=[
+        CachedResourceGroupPreparer(name_prefix="tablestest"),
+        CachedCosmosAccountPreparer(name_prefix="tablestest"),
+    ]
+)
 
 # ------------------------------------------------------------------------------
 SERVICES = {
diff --git a/tools/azure-sdk-tools/devtools_testutils/powershell_preparer.py b/tools/azure-sdk-tools/devtools_testutils/powershell_preparer.py
index abfbed4f8197..b3f8478a4b80 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 key_error
+
+                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: