From b745b139aaef97a365a29aec23cffa9c5e10d3d7 Mon Sep 17 00:00:00 2001
From: Vikram Jayanthi <vikramjayanthi@google.com>
Date: Tue, 9 Jun 2020 14:51:53 -0700
Subject: [PATCH] refactored and added tests, changed upload and get_file_size

---
 tests/test_upload.py     |  8 ++++++--
 tests/test_utils.py      | 23 ++++++++++++++---------
 twine/commands/upload.py |  7 +++----
 twine/utils.py           | 19 ++++++++++---------
 4 files changed, 33 insertions(+), 24 deletions(-)

diff --git a/tests/test_upload.py b/tests/test_upload.py
index 34770597..b30fe583 100644
--- a/tests/test_upload.py
+++ b/tests/test_upload.py
@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 import pretend
 import pytest
 import requests
@@ -19,6 +18,7 @@
 from twine import cli
 from twine import exceptions
 from twine import package as package_file
+from twine import utils
 from twine.commands import upload
 
 from . import helpers
@@ -74,7 +74,7 @@ def test_successs_prints_release_urls(upload_settings, stub_repository, capsys):
 
 
 def test_successs_prints_uploaded_package_names_if_verbose(upload_settings, capsys):
-    """Prints the path of each distribution attempting to be uploaded."""
+    """Print the path and file size of each distribution attempting to be uploaded."""
     dists_to_upload = [
         helpers.WHEEL_FIXTURE,
         helpers.SDIST_FIXTURE,
@@ -82,6 +82,8 @@ def test_successs_prints_uploaded_package_names_if_verbose(upload_settings, caps
         helpers.NEW_WHEEL_FIXTURE,
     ]
 
+    sizes_of_dists_to_upload = [utils.get_file_size(dist) for dist in dists_to_upload]
+
     upload_settings.verbose = True
 
     result = upload.upload(upload_settings, dists_to_upload)
@@ -92,6 +94,8 @@ def test_successs_prints_uploaded_package_names_if_verbose(upload_settings, caps
 
     for dist_to_upload in dists_to_upload:
         assert captured.out.count(dist_to_upload) == 1
+    for size_of_dist_to_upload in sizes_of_dists_to_upload:
+        assert captured.out.count(size_of_dist_to_upload) == 1
 
 
 def test_success_with_pre_signed_distribution(upload_settings, stub_repository):
diff --git a/tests/test_utils.py b/tests/test_utils.py
index e366fe02..b8ae34ca 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -315,12 +315,17 @@ def test_check_status_code_for_missing_status_code(capsys, repo_url):
     assert captured.out == "NOTE: Try --verbose to see response content.\n"
 
 
-def test_get_file_size_unit_determination():
-    """Test if get_file_size returns a non zero value in KB for a file < .01 MB."""
-    test_file = helpers.SDIST_FIXTURE
-
-    file_size, size_unit = utils.get_file_size(test_file)
-    print(file_size)
-    print(size_unit)
-    assert file_size != 0
-    assert size_unit == "KB"
+@pytest.mark.parametrize(
+    ("filename, expected"),
+    [
+        (helpers.WHEEL_FIXTURE, "15.4 KB"),
+        (helpers.SDIST_FIXTURE, "20.8 KB"),
+        (helpers.NEW_SDIST_FIXTURE, "26.1 KB"),
+        (helpers.NEW_WHEEL_FIXTURE, "21.9 KB"),
+    ],
+)
+def test_get_file_size(filename, expected):
+    """Compare the size of a file to the expected value."""
+    file_size = utils.get_file_size(filename)
+
+    assert file_size == expected
diff --git a/twine/commands/upload.py b/twine/commands/upload.py
index d93a37a4..ef7ff67f 100644
--- a/twine/commands/upload.py
+++ b/twine/commands/upload.py
@@ -61,10 +61,9 @@ def upload(upload_settings: settings.Settings, dists: List[str]) -> None:
 
     print(f"Uploading distributions to {repository_url}")
     if upload_settings.verbose:
-        for dist_to_upload in uploads:
-            # If the file size is <.1 MB we display in KB
-            file_size, size_unit = utils.get_file_size(dist_to_upload)
-            print(f"   {dist_to_upload} ({file_size:.1f} {size_unit})")
+        for filename in uploads:
+            file_size = utils.get_file_size(filename)
+            print(f"  {filename} ({file_size})")
         print("\n")
 
     repository = upload_settings.create_repository()
diff --git a/twine/utils.py b/twine/utils.py
index 606d44de..d859904e 100644
--- a/twine/utils.py
+++ b/twine/utils.py
@@ -23,7 +23,6 @@
 from typing import Dict
 from typing import Optional
 from typing import Sequence
-from typing import Tuple
 from typing import Union
 from urllib.parse import urlparse
 from urllib.parse import urlunparse
@@ -160,14 +159,16 @@ def normalize_repository_url(url: str) -> str:
     return urlunparse(parsed)
 
 
-def get_file_size(filepath: str) -> Tuple[float, str]:
-    """Return the file size and unit of size following PyPI's file size format."""
-    size_unit = "MB"
-    file_size = os.path.getsize(filepath) / (1024 * 1024)
-    if round(file_size, 1) == 0:
-        size_unit = "KB"
-        file_size = os.path.getsize(filepath) / 1024
-    return file_size, size_unit
+def get_file_size(filename: str) -> str:
+    """Return the size of a file, in KB if < .1 MB."""
+    file_size = os.path.getsize(filename) / 1024
+    size_unit = "KB"
+
+    if file_size > 1024:
+        file_size = file_size / 1024
+        size_unit = "MB"
+
+    return f"{file_size:.1f} {size_unit}"
 
 
 def check_status_code(response: requests.Response, verbose: bool) -> None: