Make types consistent with relay changes

getsentry · May 5, 2022 · 822ad0f · 822ad0f
1 parent 9dd1a48
commit 822ad0f
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 9 deletions.
diff --git a/sentry_sdk/_types.py b/sentry_sdk/_types.py
@@ -48,4 +48,34 @@
     ]
     SessionStatus = Literal["ok", "exited", "crashed", "abnormal"]
     EndpointType = Literal["store", "envelope"]
-    MeasurementUnit = Literal["ns", "ms", "s"]
+
+    DurationUnit = Literal[
+        "nanosecond",
+        "microsecond",
+        "millisecond",
+        "second",
+        "minute",
+        "hour",
+        "day",
+        "week",
+    ]
+
+    InformationUnit = Literal[
+        "bit",
+        "byte",
+        "kilobyte",
+        "kibibyte",
+        "megabyte",
+        "mebibyte",
+        "gigabyte",
+        "gibibyte",
+        "terabyte",
+        "tebibyte",
+        "petabyte",
+        "pebibyte",
+        "exabyte",
+        "exbibyte",
+    ]
+
+    FractionUnit = Literal["ratio", "percent"]
+    MeasurementUnit = Union[DurationUnit, InformationUnit, FractionUnit, str]
diff --git a/sentry_sdk/tracing.py b/sentry_sdk/tracing.py
@@ -605,18 +605,14 @@ def finish(self, hub=None):
 
         return hub.capture_event(event)
 
-    def set_measurement(self, name, value, unit="ms"):
+    def set_measurement(self, name, value, unit=""):
         # type: (str, float, MeasurementUnit) -> None
         if not has_custom_measurements_enabled():
             logger.debug(
                 "[Tracing] Experimental custom_measurements feature is disabled"
             )
             return
 
-        if unit not in ("ns", "ms", "s"):
-            logger.debug("[Tracing] Discarding measurement due to invalid unit")
-            return
-
         self._measurements[name] = {"value": value, "unit": unit}
 
     def to_json(self):

diff --git a/tests/tracing/test_misc.py b/tests/tracing/test_misc.py
@@ -255,9 +255,11 @@ def test_set_meaurement(sentry_init, capture_events):
 
     transaction = start_transaction(name="measuring stuff")
     transaction.set_measurement("metric.foo", 123)
-    transaction.set_measurement("metric.bar", 456, unit="s")
+    transaction.set_measurement("metric.bar", 456, unit="second")
+    transaction.set_measurement("metric.baz", 420.69, unit="custom")
     transaction.finish()
 
     (event,) = events
-    assert event["measurements"]["metric.foo"] == {"value": 123, "unit": "ms"}
-    assert event["measurements"]["metric.bar"] == {"value": 456, "unit": "s"}
+    assert event["measurements"]["metric.foo"] == {"value": 123, "unit": ""}
+    assert event["measurements"]["metric.bar"] == {"value": 456, "unit": "second"}
+    assert event["measurements"]["metric.baz"] == {"value": 420.69, "unit": "custom"}