Changing implementation and adding test for custom screenshot hook

Author: kashishm

getgauge/python.py

@@ -47,7 +47,7 @@ def after_step(obj):
 
 
 def screenshot(func):
-    registry.add_screenshot_provider(func)
+    registry.set_screenshot_provider(func)
     return func
 
 

getgauge/registry.py

@@ -44,7 +44,7 @@ def _take_screenshot():
 
 class Registry(object):
     def __init__(self):
-        self.__screenshot_providers = [_take_screenshot, _take_screenshot]
+        self.__screenshot_provider = _take_screenshot
         self.__steps_map = {}
         self.__before_step = []
         self.__after_step = []
@@ -80,7 +80,7 @@ def after_suite(self):
         return self.__after_suite
 
     def screenshot_provider(self):
-        return self.__screenshot_providers[0]
+        return self.__screenshot_provider
 
     def all_steps(self):
         return [value.step_text for value in self.__steps_map.values()]
@@ -109,8 +109,8 @@ def add_before_suite(self, func):
     def add_after_suite(self, func):
         self.__after_suite.append(func)
 
-    def add_screenshot_provider(self, func):
-        self.__screenshot_providers[0] = func
+    def set_screenshot_provider(self, func):
+        self.__screenshot_provider = func
 
     def add_step_definition(self, step_text, func, file_name, has_alias=False):
         if not isinstance(step_text, list):

tests/test_data/impl_stub.py

@@ -1,5 +1,5 @@
 from getgauge.python import before_step, step, after_step, before_scenario, after_scenario, before_spec, after_spec, \
-    before_suite, after_suite
+    before_suite, after_suite, screenshot
 
 
 @step("Step 1")
@@ -80,3 +80,7 @@ def after_suite1():
 @after_suite
 def after_suite2():
     pass
+
+@screenshot
+def take_screenshot():
+    return ""

tests/test_gauge.py tests/test_python.pytests/test_gauge.py renamed to tests/test_python.py

@@ -306,6 +306,11 @@ def test_after_suite_decorator(self):
         expected = ['after_suite1', 'after_suite2']
         self.assertEqual(expected, [func.func_name for func in funcs])
 
+    def test_screenshot_decorator(self):
+        func = registry.screenshot_provider()
+        expected = 'take_screenshot'
+        self.assertEqual(expected, func.func_name)
+
 
 if __name__ == '__main__':
     unittest.main()