Merge pull request #494 from yuvipanda/util-tests

Add unit tests for repo2docker/utils.py

Author: betatim

repo2docker/app.py

@@ -14,6 +14,7 @@
 import os
 import pwd
 import subprocess
+import shutil
 import tempfile
 import time
 
@@ -35,7 +36,7 @@
 )
 from . import contentproviders
 from .utils import (
-    ByteSpecification, maybe_cleanup, is_valid_docker_image_name,
+    ByteSpecification, is_valid_docker_image_name,
     validate_and_generate_port_mapping, chdir
 )
 
@@ -723,9 +724,7 @@ def start(self):
             else:
                 checkout_path = self.git_workdir
 
-        # keep as much as possible in the context manager to make sure we
-        # cleanup if things go wrong
-        with maybe_cleanup(checkout_path, self.cleanup_checkout):
+        try:
             self.fetch(self.repo, self.ref, checkout_path)
 
             if self.subdir:
@@ -771,6 +770,10 @@ def start(self):
                         else:
                             self.log.info(json.dumps(l),
                                           extra=dict(phase='building'))
+        finally:
+            # Cheanup checkout if necessary
+            if self.cleanup_checkout:
+                shutil.rmtree(checkout_path, ignore_errors=True)
 
         if self.push:
             self.push_image()

repo2docker/utils.py

@@ -2,15 +2,16 @@
 from functools import partial
 import os
 import re
-import shutil
 import subprocess
 
 from traitlets import Integer, TraitError
 
 
 def execute_cmd(cmd, capture=False, **kwargs):
     """
-    Call given command, yielding output line by line if capture=True
+    Call given command, yielding output line by line if capture=True.
+
+    Must be yielded from.
     """
     if capture:
         kwargs['stdout'] = subprocess.PIPE
@@ -67,14 +68,6 @@ def chdir(path):
         os.chdir(old_dir)
 
 
-@contextmanager
-def maybe_cleanup(path, cleanup=False):
-    """Delete the directory at passed path if cleanup flag is True."""
-    yield
-    if cleanup:
-        shutil.rmtree(path, ignore_errors=True)
-
-
 def validate_and_generate_port_mapping(port_mapping):
     """
     Validate the port mapping list and return a list of validated tuples.

tests/test_utils.py

@@ -0,0 +1,66 @@
+"""
+Tests for repo2docker/utils.py
+"""
+import traitlets
+import os
+from tempfile import TemporaryDirectory
+from repo2docker import utils
+import pytest
+import subprocess
+
+
+def test_capture_cmd_no_capture_success():
+    # This should succeed
+    for line in utils.execute_cmd([
+        '/bin/bash', '-c', 'echo test'
+    ]):
+        pass
+
+def test_capture_cmd_no_capture_fail():
+    with pytest.raises(subprocess.CalledProcessError):
+        for line in utils.execute_cmd([
+            '/bin/bash', '-c', 'e '
+        ]):
+            pass
+
+
+def test_capture_cmd_capture_success():
+    # This should succeed
+    for line in utils.execute_cmd([
+        '/bin/bash', '-c', 'echo test'
+    ], capture=True):
+        assert line == 'test\n'
+
+
+def test_capture_cmd_capture_fail():
+    with pytest.raises(subprocess.CalledProcessError):
+        for line in utils.execute_cmd([
+            '/bin/bash', '-c', 'echo test; exit 1 '
+        ], capture=True):
+            assert line == 'test\n'
+
+
+def test_chdir():
+    with TemporaryDirectory() as d:
+        cur_cwd = os.getcwd()
+        with utils.chdir(d):
+            assert os.getcwd() == d
+        assert os.getcwd() == cur_cwd
+
+
+def test_byte_spec_validation():
+    bs = utils.ByteSpecification()
+
+    assert bs.validate(None, 1) == 1
+    assert bs.validate(None, 1.0) == 1.0
+
+    assert bs.validate(None, '1K') == 1024
+    assert bs.validate(None, '1M') == 1024 * 1024
+    assert bs.validate(None, '1G') == 1024 * 1024 * 1024
+    assert bs.validate(None, '1T') == 1024 * 1024 * 1024 * 1024
+
+    with pytest.raises(traitlets.TraitError):
+        bs.validate(None, 'NK')
+
+    with pytest.raises(traitlets.TraitError):
+        bs.validate(None, '1m')