[Piglit] [PATCH 10/14] framework/core: extend check_dir to take a handler callable

[Dylan Baker]

From: Dylan Baker <baker.dylan.c at gmail.com>

This allows check_dir to fulfill the needs of the overwrite switch in
the main run function.

Signed-off-by: Dylan Baker <dylanx.c.baker at intel.com>
---
 framework/core.py         | 18 ++++++++++++++----
 framework/programs/run.py | 27 ++++++++++++++++-----------
 unittests/core_tests.py   |  9 +++++++++
 3 files changed, 39 insertions(+), 15 deletions(-)

diff --git a/framework/core.py b/framework/core.py
index 60213c2..6832c0a 100644
--- a/framework/core.py
+++ b/framework/core.py
@@ -114,8 +114,8 @@ def get_config(arg=None):
                 pass
 
 
-def check_dir(dirname, failifexists=False):
-    """Check for the existance of a directory and create it if possible.
+def check_dir(dirname, failifexists=False, handler=None):
+    """Check for the existence of a directory and create it if possible.
 
     This function will check for the existance of a directory. If that
     directory doesn't exist it will try to create it. If the directory does
@@ -124,21 +124,31 @@ def check_dir(dirname, failifexists=False):
     2) If "failifexists" is True it will raise an PiglitException, it is the
     job of the caller using failifexists=True to handle this exception
 
+    Both failifexists and handler can be passed, but failifexists will have
+    precedence.
+
     Arguments:
     dirname -- the name of the directory to check
 
     Keyword Arguments:
     failifexists -- If True and the directory exists then PiglitException will
                     be raised (default: False)
+    handler -- a callable that is passed dirname if the thing to check exists.
 
     """
     try:
         os.stat(dirname)
     except OSError as e:
-        if e.errno not in [errno.ENOENT, errno.ENOTDIR] and failifexists:
-            raise exceptions.PiglitException
+        # If the error is not "no file or directory" or "not a dir", then
+        # either raise an exception, call the handler function, or return
+        if e.errno not in [errno.ENOENT, errno.ENOTDIR]:
+            if failifexists:
+                raise exceptions.PiglitException
+            elif handler is not None:
+                handler(dirname)
 
     try:
+        # makedirs is expensive, so check before # calling it.
         if not os.path.exists(dirname):
             os.makedirs(dirname)
     except OSError as e:
diff --git a/framework/programs/run.py b/framework/programs/run.py
index dff9f49..f1b237d 100644
--- a/framework/programs/run.py
+++ b/framework/programs/run.py
@@ -212,6 +212,14 @@ def _disable_windows_exception_messages():
         ctypes.windll.kernel32.SetErrorMode(uMode)
 
 
+def _results_handler(path):
+    """Handler for core.check_dir."""
+    if os.path.isdir(path):
+        shutil.rmtree(path)
+    else:
+        os.unlink(path)
+
+
 @exceptions.handler
 def run(input_):
     """ Function for piglit run command
@@ -246,17 +254,14 @@ def run(input_):
 
     # If the results directory already exists and if overwrite was set, then
     # clear the directory. If it wasn't set, then raise fatal error.
-    if os.path.exists(args.results_path):
-        if args.overwrite:
-            if os.path.isdir(args.results_path):
-                shutil.rmtree(args.results_path)
-            else:
-                os.unlink(args.results_path)
-        else:
-            raise exceptions.PiglitFatalError(
-                'Cannot overwrite existing folder without the -o/--overwrite '
-                'option being set.')
-    os.makedirs(args.results_path)
+    try:
+        core.check_dir(args.results_path,
+                       failifexists=args.overwrite,
+                       handler=_results_handler)
+    except exceptions.PiglitException:
+        raise exceptions.PiglitFatalError(
+            'Cannot overwrite existing folder without the -o/--overwrite '
+            'option being set.')
 
     results = framework.results.TestrunResult()
     backends.set_meta(args.backend, results)
diff --git a/unittests/core_tests.py b/unittests/core_tests.py
index c0ea523..8cbac8e 100644
--- a/unittests/core_tests.py
+++ b/unittests/core_tests.py
@@ -321,3 +321,12 @@ def test_check_dir_makedirs_fail():
             with mock.patch('framework.core.os.makedirs',
                             mock.Mock(side_effect=OSError)):
                 core.check_dir('foo', False)
+
+
+ at nt.raises(utils.SentinalException)
+def test_check_dir_handler():
+    """core.check_dir: Handler is called if not failifexists."""
+    with mock.patch('framework.core.os.stat',
+                    mock.Mock(side_effect=OSError('foo', errno.ENOTDIR))):
+        core.check_dir('foo',
+                       handler=mock.Mock(side_effect=utils.SentinalException))
-- 
2.8.2