From bfd921bdbc690fa74eba30cb58b81b8e3f5520c3 Mon Sep 17 00:00:00 2001
From: Jeremy Stretch <jstretch@netboxlabs.com>
Date: Thu, 6 Jun 2024 09:24:29 -0400
Subject: [PATCH] Rewrite sizing logic

---
 netbox/templates/virtualization/cluster.html  |  2 +-
 .../virtualization/virtualmachine.html        |  2 +-
 netbox/utilities/templatetags/helpers.py      | 35 ++++++-------------
 .../tests/test_humanize_megabytes_helper.py   | 25 -------------
 4 files changed, 13 insertions(+), 51 deletions(-)
 delete mode 100644 netbox/utilities/tests/test_humanize_megabytes_helper.py

diff --git a/netbox/templates/virtualization/cluster.html b/netbox/templates/virtualization/cluster.html
index c1ac74603..a2c3e06c1 100644
--- a/netbox/templates/virtualization/cluster.html
+++ b/netbox/templates/virtualization/cluster.html
@@ -59,7 +59,7 @@
               <th scope="row"><i class="mdi mdi-chip"></i> {% trans "Memory" %}</th>
               <td>
                   {% if memory_sum %}
-                    <span title={{ memory_sum }}>{{ memory_sum|humanize_megabytes }}</span>
+                      <span title={{ memory_sum }}>{{ memory_sum|humanize_megabytes }}</span>
                   {% else %}
                       {{ ''|placeholder }}
                   {% endif %}
diff --git a/netbox/templates/virtualization/virtualmachine.html b/netbox/templates/virtualization/virtualmachine.html
index 2d87c6df0..ed8980f5c 100644
--- a/netbox/templates/virtualization/virtualmachine.html
+++ b/netbox/templates/virtualization/virtualmachine.html
@@ -125,7 +125,7 @@
                     <th scope="row"><i class="mdi mdi-chip"></i> {% trans "Memory" %}</th>
                     <td>
                         {% if object.memory %}
-                          <span title={{ object.memory }}> {{ object.memory|humanize_megabytes }}</span>
+                            <span title={{ object.memory }}>{{ object.memory|humanize_megabytes }}</span>
                         {% else %}
                             {{ ''|placeholder }}
                         {% endif %}
diff --git a/netbox/utilities/templatetags/helpers.py b/netbox/utilities/templatetags/helpers.py
index 4f08dc107..b9a3e0005 100644
--- a/netbox/utilities/templatetags/helpers.py
+++ b/netbox/utilities/templatetags/helpers.py
@@ -1,14 +1,9 @@
-import datetime
 import json
 from typing import Dict, Any
 from urllib.parse import quote
 
 from django import template
-from django.conf import settings
-from django.template.defaultfilters import date
 from django.urls import NoReverseMatch, reverse
-from django.utils import timezone
-from django.utils.safestring import mark_safe
 
 from core.models import ObjectType
 from utilities.forms import get_selected_values, TableConfigForm
@@ -92,30 +87,22 @@ def humanize_speed(speed):
 @register.filter()
 def humanize_megabytes(mb):
     """
-    Express a number of megabytes in the most suitable unit (e.g. gigabytes or terabytes).
-    It considers the mb value as megabytes and converts it to the most suitable unit.
+    Express a number of megabytes in the most suitable unit (e.g. gigabytes, terabytes, etc.).
     """
-
-    # Factors in bytes
-    factors = {
-        "mega": 1000 ** 2,
-        "giga": 1000 ** 3,
-        "tera": 1000 ** 4,
-    }
-
     if not mb:
         return ""
 
-    bytes = int(mb * 1000**2)
+    PB_SIZE = 1000000000
+    TB_SIZE = 1000000
+    GB_SIZE = 1000
 
-    if bytes >= factors["tera"]:
-        return f"{bytes / factors['tera']:.2f} TB"
-
-    if bytes >= factors["giga"]:
-        return f"{bytes / factors['giga']:.2f} GB"
-
-    if bytes >= factors["mega"]:
-        return f"{bytes / factors['mega']:.2f} MB"
+    if mb >= PB_SIZE:
+        return f"{mb / PB_SIZE:.2f} PB"
+    if mb >= TB_SIZE:
+        return f"{mb / TB_SIZE:.2f} TB"
+    if mb >= GB_SIZE:
+        return f"{mb / GB_SIZE:.2f} GB"
+    return f"{mb} MB"
 
 
 @register.filter()
diff --git a/netbox/utilities/tests/test_humanize_megabytes_helper.py b/netbox/utilities/tests/test_humanize_megabytes_helper.py
deleted file mode 100644
index 642e92b9f..000000000
--- a/netbox/utilities/tests/test_humanize_megabytes_helper.py
+++ /dev/null
@@ -1,25 +0,0 @@
-from utilities.templatetags.helpers import humanize_megabytes
-from utilities.testing import TestCase
-
-
-class TestConvertByteSize(TestCase):
-
-    def test_humanize_megabytes_converts_megabytes(self):
-        """Test that humanize_megabytes converts megabytes to the most suitable unit."""
-        self.assertEqual(humanize_megabytes(1), "1.00 MB")
-
-    def test_humanize_megabytes_converts_to_gigabytes(self):
-        """Test that humanize_megabytes converts megabytes to gigabytes."""
-        self.assertEqual(humanize_megabytes(1000), "1.00 GB")
-
-    def test_humanize_megabytes_converts_to_terabytes(self):
-        """Test that humanize_megabytes converts megabytes to terabytes."""
-        self.assertEqual(humanize_megabytes(1000000), "1.00 TB")
-
-    def test_humanize_megabytes_returns_empty_for_none(self):
-        """Test that humanize_megabytes returns empty for None."""
-        self.assertEqual(humanize_megabytes(None), '')
-
-    def test_humanize_megabytes_without_unit(self):
-        """Test that humanize_megabytes returns the value without unit."""
-        self.assertEqual(humanize_megabytes(123456789), "123.46 TB")