fix test cases

netbox/utilities/forms/fields.py

@@ -17,7 +17,7 @@ from django.urls import reverse
 from utilities.choices import unpack_grouped_choices
 from utilities.utils import content_type_identifier, content_type_name
 from utilities.validators import EnhancedURLValidator
-from virtualization.choices import MemoryUnitChoices
+from virtualization.choices import *
 from . import widgets
 from .constants import *
 from .utils import expand_alphanumeric_pattern, expand_ipaddress_pattern, parse_csv, validate_csv
@@ -35,7 +35,6 @@ __all__ = (
     'CSVMultipleChoiceField',
     'CSVMultipleContentTypeField',
     'CSVTypedChoiceField',
-    'DiskField',
     'DynamicModelChoiceField',
     'DynamicModelMultipleChoiceField',
     'ExpandableIPAddressField',
@@ -154,68 +153,36 @@ class MACAddressField(forms.Field):
 
 
 class MemoryField(forms.MultiValueField):
-    widget = widgets.MemoryWidget
+    """
-    MULTIPLIERS = {
+    Memory Unit Field
-        MemoryUnitChoices.UNIT_MB: 1024**0,
+    """
-        MemoryUnitChoices.UNIT_GB: 1024**1,
-        MemoryUnitChoices.UNIT_TB: 1024**2,
-    }
 
     def __init__(self, **kwargs):
+        self.multipliers = kwargs['multipliers']
+
+        widget = widgets.MemoryWidget(
+            choices=kwargs['choices'],
+            default_unit=kwargs['default_unit']
+        )
         fields = (
             forms.IntegerField(required=False),
             forms.ChoiceField(
-                choices=MemoryUnitChoices.MEMORY_CHOICES,
+                choices=kwargs['choices'],
                 required=False
             ),
         )
         super(MemoryField, self).__init__(
             fields=fields, required=False,
-            require_all_fields=False, **kwargs
+            require_all_fields=False, widget=widget
         )
 
     @classmethod
     def compress(cls, data):
         if data:
             size, unit = data
-            if size and not unit:
+            if not size or not unit:
-                raise forms.ValidationError("Memory unit cannot be blank.")
+                raise forms.ValidationError("Size and unit must be specified.")
-            elif size and unit:
+            return size * cls.multipliers[unit]
-                return size * cls.MULTIPLIERS[unit]
-            elif not size and unit:
-                raise forms.ValidationError("Please enter a memory value when unit is selected.")
-
-
-class DiskField(forms.MultiValueField):
-    widget = widgets.DiskWidget
-    MULTIPLIERS = {
-        MemoryUnitChoices.UNIT_GB: 1024**0,
-        MemoryUnitChoices.UNIT_TB: 1024**1,
-    }
-
-    def __init__(self, **kwargs):
-        fields = (
-            forms.IntegerField(required=False),
-            forms.ChoiceField(
-                choices=MemoryUnitChoices.DISK_CHOICES,
-                required=False
-            ),
-        )
-        super(DiskField, self).__init__(
-            fields=fields, required=False,
-            require_all_fields=False, **kwargs
-        )
-
-    @classmethod
-    def compress(cls, data):
-        if data:
-            size, unit = data
-            if size and not unit:
-                raise forms.ValidationError("Disk unit cannot be blank.")
-            elif size and unit:
-                return size * cls.MULTIPLIERS[unit]
-            elif not size and unit:
-                raise forms.ValidationError("Please enter a disk value when unit is selected.")
 
 
 #

netbox/utilities/forms/widgets.py

@@ -7,7 +7,7 @@ from django.contrib.postgres.forms import SimpleArrayField
 
 from utilities.choices import ColorChoices
 from .utils import add_blank_choice, parse_numeric_range
-from virtualization.choices import MemoryUnitChoices
+from virtualization.choices import UNIT_GB, UNIT_MB, UNIT_TB
 
 __all__ = (
     'APISelect',
@@ -17,7 +17,6 @@ __all__ = (
     'ColorSelect',
     'DatePicker',
     'DateTimePicker',
-    'DiskWidget',
     'MemoryWidget',
     'NumericArrayField',
     'SelectSpeedWidget',
@@ -317,35 +316,18 @@ class TimePicker(forms.TextInput):
 
 class MemoryWidget(forms.MultiWidget):
     """
-    Memory Widget.
+    Memory Unit Widget.
     """
-    def __init__(self, attrs=None):
+    def __init__(self, attrs=None, **kwargs):
+        self.default_unit = kwargs['default_unit']
         widgets = (
             forms.NumberInput(),
-            StaticSelect(choices=add_blank_choice(MemoryUnitChoices.MEMORY_CHOICES))
+            StaticSelect(choices=add_blank_choice(kwargs['choices']))
         )
         super(MemoryWidget, self).__init__(widgets, attrs)
 
     def decompress(self, value):
         if value:
-            return [value, MemoryUnitChoices.UNIT_MB]
+            return [value, self.default_unit]
-        else:
-            return ['', '']
-
-
-class DiskWidget(forms.MultiWidget):
-    """
-    Disk Widget.
-    """
-    def __init__(self, attrs=None):
-        widgets = (
-            forms.NumberInput(),
-            StaticSelect(choices=add_blank_choice(MemoryUnitChoices.DISK_CHOICES))
-        )
-        super(DiskWidget, self).__init__(widgets, attrs)
-
-    def decompress(self, value):
-        if value:
-            return [value, MemoryUnitChoices.UNIT_GB]
         else:
             return ['', '']

netbox/virtualization/choices.py

@@ -5,17 +5,33 @@ from utilities.choices import ChoiceSet
 # VirtualMachines
 #
 
-class MemoryUnitChoices(ChoiceSet):
+UNIT_MB = 'mb'
-    UNIT_MB = 'mb'
+UNIT_GB = 'gb'
-    UNIT_GB = 'gb'
+UNIT_TB = 'tb'
-    UNIT_TB = 'tb'
 
-    DISK_CHOICES = (
+MEMORY_MULTIPLIERS = {
+    UNIT_MB: 1024**0,
+    UNIT_GB: 1024**1,
+    UNIT_TB: 1024**2,
+}
+
+DISK_MULTIPLIERS = {
+    UNIT_GB: 1024**0,
+    UNIT_TB: 1024**1,
+}
+
+
+class DiskUnitChoices(ChoiceSet):
+
+    CHOICES = (
         (UNIT_GB, 'GB'),
         (UNIT_TB, 'TB'),
     )
 
-    MEMORY_CHOICES = (
+
+class MemoryUnitChoices(ChoiceSet):
+
+    CHOICES = (
         (UNIT_MB, 'MB'),
         (UNIT_GB, 'GB'),
         (UNIT_TB, 'TB'),

netbox/virtualization/forms/bulk_edit.py

@@ -7,7 +7,7 @@ from extras.forms import AddRemoveTagsForm, CustomFieldModelBulkEditForm
 from ipam.models import VLAN
 from tenancy.models import Tenant
 from utilities.forms import (
-    add_blank_choice, BulkEditNullBooleanSelect, BulkRenameForm, CommentField, DiskField, DynamicModelChoiceField,
+    add_blank_choice, BulkEditNullBooleanSelect, BulkRenameForm, CommentField, DynamicModelChoiceField,
     DynamicModelMultipleChoiceField, MemoryField, SmallTextarea, StaticSelect
 )
 from virtualization.choices import *
@@ -131,8 +131,16 @@ class VirtualMachineBulkEditForm(AddRemoveTagsForm, CustomFieldModelBulkEditForm
         required=False,
         label='vCPUs'
     )
-    memory = MemoryField()
+    memory = MemoryField(
-    disk = DiskField()
+        multipliers=MEMORY_MULTIPLIERS,
+        choices=MemoryUnitChoices.CHOICES,
+        default_unit=UNIT_MB
+    )
+    disk = MemoryField(
+        multipliers=DISK_MULTIPLIERS,
+        choices=DiskUnitChoices.CHOICES,
+        default_unit=UNIT_MB
+    )
     comments = CommentField(
         widget=SmallTextarea,
         label='Comments'

netbox/virtualization/forms/models.py

@@ -10,10 +10,11 @@ from extras.models import Tag
 from ipam.models import IPAddress, VLAN, VLANGroup
 from tenancy.forms import TenancyForm
 from utilities.forms import (
-    BootstrapMixin, CommentField, ConfirmationForm, DiskField, DynamicModelChoiceField, DynamicModelMultipleChoiceField,
+    BootstrapMixin, CommentField, ConfirmationForm, DynamicModelChoiceField, DynamicModelMultipleChoiceField,
     JSONField, MemoryField, SlugField, StaticSelect
 )
 from virtualization.models import *
+from virtualization.choices import *
 
 __all__ = (
     'ClusterAddDevicesForm',
@@ -205,8 +206,8 @@ class VirtualMachineForm(TenancyForm, CustomFieldModelForm):
         queryset=Tag.objects.all(),
         required=False
     )
-    memory = MemoryField()
+    memory = MemoryField(multipliers=MEMORY_MULTIPLIERS, choices=MemoryUnitChoices.CHOICES, default_unit=UNIT_MB)
-    disk = DiskField()
+    disk = MemoryField(multipliers=DISK_MULTIPLIERS, choices=DiskUnitChoices.CHOICES, default_unit=UNIT_GB)
 
     class Meta:
         model = VirtualMachine

netbox/virtualization/tests/test_views.py

@@ -193,8 +193,9 @@ class VirtualMachineTestCase(ViewTestCases.PrimaryObjectViewTestCase):
             'primary_ip6': None,
             'vcpus': 4,
             'memory_0': 32768,
-            'memory_1': MemoryUnitChoices.UNIT_MB,
+            'memory_1': UNIT_MB,
-            'disk': 4000,
+            'disk_0': 4000,
+            'disk_1': UNIT_GB,
             'comments': 'Some comments',
             'tags': [t.pk for t in tags],
             'local_context_data': None,
@@ -214,8 +215,10 @@ class VirtualMachineTestCase(ViewTestCases.PrimaryObjectViewTestCase):
             'status': VirtualMachineStatusChoices.STATUS_STAGED,
             'role': deviceroles[1].pk,
             'vcpus': 8,
-            'memory': 65535,
+            'memory_0': 65535,
-            'disk': 8000,
+            'memory_1': UNIT_MB,
+            'disk_0': 8000,
+            'disk_1': UNIT_GB,
             'comments': 'New comments',
         }