From 8defe1e49762f63aad96d2cc01cdc59a9d59b6c1 Mon Sep 17 00:00:00 2001
From: Ivan Devat <idevat@redhat.com>
Date: Tue, 16 Jan 2018 14:23:22 +0100
Subject: [PATCH] allow to make a local config. shortcut in tests

---
 pcs/test/tools/command_env/config.py | 11 +++++++++++
 pcs/test/tools/command_env/tools.py  | 12 ++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/pcs/test/tools/command_env/config.py b/pcs/test/tools/command_env/config.py
index 4ec5b4bc7..b0e51b0ab 100644
--- a/pcs/test/tools/command_env/config.py
+++ b/pcs/test/tools/command_env/config.py
@@ -31,6 +31,17 @@ def __init__(self):
 
         self.spy = None
 
+    def add_extension(self, name, Extension):
+        if hasattr(self, name):
+            raise AssertionError(
+                "Config (integration tests) has the extension '{0}' already."
+                .format(name)
+            )
+        setattr(self, name, self.__wrap_helper(
+            Extension(self.__calls, self.__wrap_helper, self)
+        ))
+
+
     def set_spy(self, auth_tokens, ports=None):
         self.spy = Spy(auth_tokens, ports)
         return self
diff --git a/pcs/test/tools/command_env/tools.py b/pcs/test/tools/command_env/tools.py
index 2331d419b..6ed0198e6 100644
--- a/pcs/test/tools/command_env/tools.py
+++ b/pcs/test/tools/command_env/tools.py
@@ -18,12 +18,17 @@ def get_env_tools(
     default_wait_timeout=DEFAULT_WAIT_TIMEOUT,
     default_wait_error_returncode=WAIT_TIMEOUT_EXPIRED_RETURNCODE,
     exception_reports_in_processor_by_default=True,
+    local_extensions=None
 ):
     """
     Shortcut for preparing EnvAssistant and Config
 
     TestCase test_case -- corresponding test_case is used to registering cleanup
         method - to assert that everything is finished
+    dict local_extensions -- key is name of local extension, value is class that
+        will be used for local extension. So in config will be possible use
+        something like this:
+            config.my_local_extension.my_local_call_shortcut()
     """
 
     env_assistant = EnvAssistant(
@@ -38,4 +43,11 @@ def get_env_tools(
     runner.pcmk.default_wait_timeout = default_wait_timeout
     runner.pcmk.default_wait_error_returncode = default_wait_error_returncode
 
+    if local_extensions:
+        for name, ExtensionClass in local_extensions.items():
+            env_assistant.config.add_extension(
+                name,
+                ExtensionClass,
+            )
+
     return env_assistant, env_assistant.config