From 822ad0f4148a55eba95e4f4564fb2735e8bd0f18 Mon Sep 17 00:00:00 2001
From: Neel Shah <neelshah.sa@gmail.com>
Date: Thu, 5 May 2022 14:31:12 +0200
Subject: [PATCH] Make types consistent with relay changes

---
 sentry_sdk/_types.py       | 32 +++++++++++++++++++++++++++++++-
 sentry_sdk/tracing.py      |  6 +-----
 tests/tracing/test_misc.py |  8 +++++---
 3 files changed, 37 insertions(+), 9 deletions(-)

diff --git a/sentry_sdk/_types.py b/sentry_sdk/_types.py
index 6e66ad91ef..59970ad60a 100644
--- 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
index 45cf98d5a5..17b94be90f 100644
--- a/sentry_sdk/tracing.py
+++ b/sentry_sdk/tracing.py
@@ -605,7 +605,7 @@ 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(
@@ -613,10 +613,6 @@ def set_measurement(self, name, value, unit="ms"):
             )
             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
index 3b4102c294..dd47f25220 100644
--- 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"}