2752: Factor out arv-put byte tally, and test.

diff --git a/sdk/python/arvados/commands/put.py b/sdk/python/arvados/commands/put.py
index d3da8cfc8c687819c4714aae5ff1c7f332396292..5de86168d0bb8bfa005e58f4d08ec1903f4bd7e8 100644 (file)
--- a/sdk/python/arvados/commands/put.py
+++ b/sdk/python/arvados/commands/put.py
@@ -263,6 +263,18 @@ class CollectionWriterWithProgress(arvados.CollectionWriter):
         return manifest_text
 
 
+def expected_bytes_for(pathlist):
+    bytesum = 0
+    for path in pathlist:
+        if os.path.isdir(path):
+            for filename in arvados.util.listdir_recursive(path):
+                bytesum += os.path.getsize(os.path.join(path, filename))
+        elif not os.path.isfile(path):
+            return None
+        else:
+            bytesum += os.path.getsize(path)
+    return bytesum
+
 def main(arguments=None):
     args = parse_arguments(arguments)
 
@@ -277,17 +289,9 @@ def main(arguments=None):
 
     # Walk the given directory trees and stat files, adding up file sizes,
     # so we can display progress as percent
-    writer.bytes_expected = 0
-    for path in args.paths:
-        if os.path.isdir(path):
-            for filename in arvados.util.listdir_recursive(path):
-                writer.bytes_expected += os.path.getsize(
-                    os.path.join(path, filename))
-        elif not os.path.isfile(path):
-            del writer.bytes_expected
-            break
-        else:
-            writer.bytes_expected += os.path.getsize(path)
+    writer.bytes_expected = expected_bytes_for(args.paths)
+    if writer.bytes_expected is None:
+        del writer.bytes_expected
 
     # Copy file data to Keep.
     for path in args.paths:

diff --git a/sdk/python/tests/test_arv-put.py b/sdk/python/tests/test_arv-put.py
index 8412c6a34a1abc3202dcd5beba6272f080dca7e2..50060efb82226f19e8520ece7312a0a54c5d37eb 100644 (file)
--- a/sdk/python/tests/test_arv-put.py
+++ b/sdk/python/tests/test_arv-put.py
@@ -3,6 +3,7 @@
 
 import os
 import re
+import shutil
 import tempfile
 import unittest
 
@@ -227,6 +228,27 @@ class ArvadosPutResumeCacheCollectionWriterTest(ArvadosKeepLocalStoreTestCase):
         self.assertEquals(cwriter.manifest_text(), new_writer.manifest_text())
 
 
+class ArvadosExpectedBytesTest(ArvadosBaseTestCase):
+    TEST_SIZE = os.path.getsize(__file__)
+
+    def test_expected_bytes_for_file(self):
+        self.assertEquals(self.TEST_SIZE,
+                          arv_put.expected_bytes_for([__file__]))
+
+    def test_expected_bytes_for_tree(self):
+        tree = self.make_tmpdir()
+        shutil.copyfile(__file__, os.path.join(tree, 'one'))
+        shutil.copyfile(__file__, os.path.join(tree, 'two'))
+        self.assertEquals(self.TEST_SIZE * 2,
+                          arv_put.expected_bytes_for([tree]))
+        self.assertEquals(self.TEST_SIZE * 3,
+                          arv_put.expected_bytes_for([tree, __file__]))
+
+    def test_expected_bytes_for_device(self):
+        self.assertIsNone(arv_put.expected_bytes_for(['/dev/null']))
+        self.assertIsNone(arv_put.expected_bytes_for([__file__, '/dev/null']))
+
+
 class ArvadosPutTest(ArvadosKeepLocalStoreTestCase):
     def test_simple_file_put(self):
         with self.make_test_file() as testfile: