change per @penaguerrero comment

src/stcal/testing_helpers.py

@@ -1,5 +1,7 @@
 import tracemalloc
 
+MEMORY_UNIT_CONVERSION = {"B": 1, "KB": 1024, "MB": 1024 ** 2, "GB": 1024 ** 3}
+
 class MemoryThresholdExceeded(Exception):
     pass
 
@@ -13,14 +15,19 @@ class MemoryThreshold:
         # code that should not exceed expected
     """
 
-    def __init__(self, expected_usage):
+    def __init__(self, expected_usage, log_units="KB"):
         """
         Parameters
         ----------
         expected_usage : int
             Expected peak memory usage in bytes
+
+        log_units : str, optional
+            Units in which to display memory usage for error message. 
+            Supported are "B", "KB", "MB", "GB". Default is "KB".
         """
         self.expected_usage = expected_usage
+        self.log_units = log_units
 
     def __enter__(self):
         tracemalloc.start()
@@ -31,6 +38,8 @@ def __exit__(self, exc_type, exc_value, traceback):
         tracemalloc.stop()
 
         if peak > self.expected_usage:
+            scaling = MEMORY_UNIT_CONVERSION[self.log_units]
             msg = ("Peak memory usage exceeded expected usage: "
-                  f"{peak / 1024:.2f} KB > {self.expected_usage / 1024:.2f} KB")
+                  f"{peak / scaling:.2f} {self.log_units} > "
+                  f"{self.expected_usage / scaling:.2f} {self.log_units} ")
             raise MemoryThresholdExceeded(msg)

tests/test_infrastructure.py

@@ -10,7 +10,7 @@ def test_memory_threshold():
         buff = np.empty(200, dtype=np.uint8)
 
 
-def test_memory_threshold_raise():
+def test_memory_threshold_exceeded():
     with pytest.raises(MemoryThresholdExceeded):
         with MemoryThreshold(1000):
             buff = np.empty(2000, dtype=np.uint8)