From 8e30d7ded1f6679504cb61879226c853d71fdc87 Mon Sep 17 00:00:00 2001
From: Leighton Chen <lechen@microsoft.com>
Date: Thu, 30 Jun 2022 10:36:58 -0700
Subject: [PATCH] Add kwargs to derived gauge (#1135)

---
 CHANGELOG.md                              |  5 +++--
 contrib/opencensus-ext-azure/CHANGELOG.md |  2 ++
 opencensus/metrics/export/gauge.py        | 18 ++++++++++--------
 tests/unit/metrics/export/test_gauge.py   |  9 +++++++++
 4 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index cbec5c1db..8b7a6e90e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,8 +1,9 @@
 # Changelog
 
 ## Unreleased
-- Allow specifying metrics (custom_measurements) for Azure custom events
-([#1117](https://github.com/census-instrumentation/opencensus-python/pull/1117))
+
+- Add kwargs to derived gauge
+([#1135](https://github.com/census-instrumentation/opencensus-python/pull/1135))
 
 # 0.9.0
 Released 2022-04-20
diff --git a/contrib/opencensus-ext-azure/CHANGELOG.md b/contrib/opencensus-ext-azure/CHANGELOG.md
index e70c79ab1..d84f3e28f 100644
--- a/contrib/opencensus-ext-azure/CHANGELOG.md
+++ b/contrib/opencensus-ext-azure/CHANGELOG.md
@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Allow specifying metrics (custom_measurements) for Azure custom events
+([#1117](https://github.com/census-instrumentation/opencensus-python/pull/1117))
 - Shutdown Statsbeat when hitting error/exception threshold
 ([#1127](https://github.com/census-instrumentation/opencensus-python/pull/1127))
 - Fix failure counting statsbeat - refactor status code logic in transport
diff --git a/opencensus/metrics/export/gauge.py b/opencensus/metrics/export/gauge.py
index 10b35419e..5ae05985c 100644
--- a/opencensus/metrics/export/gauge.py
+++ b/opencensus/metrics/export/gauge.py
@@ -194,15 +194,17 @@ class DerivedGaugePoint(GaugePoint):
         :class:`opencensus.metrics.export.cumulative.CumulativePointDouble`
     :param gauge_point: The underlying `GaugePoint`.
     """
-    def __init__(self, func, gauge_point):
+    def __init__(self, func, gauge_point, **kwargs):
         self.gauge_point = gauge_point
         self.func = utils.get_weakref(func)
+        self._kwargs = kwargs
 
     def __repr__(self):
-        return ("{}({})"
+        return ("{}({})({})"
                 .format(
                     type(self).__name__,
-                    self.func()
+                    self.func(),
+                    self._kwargs
                 ))
 
     def get_value(self):
@@ -216,7 +218,7 @@ def get_value(self):
             longer exists.
         """
         try:
-            val = self.func()()
+            val = self.func()(**self._kwargs)
         except TypeError:  # The underlying function has been GC'd
             return None
 
@@ -406,13 +408,13 @@ class DerivedGauge(BaseGauge):
     instead of using this class directly.
     """
 
-    def _create_time_series(self, label_values, func):
+    def _create_time_series(self, label_values, func, **kwargs):
         with self._points_lock:
             return self.points.setdefault(
                 tuple(label_values),
-                DerivedGaugePoint(func, self.point_type()))
+                DerivedGaugePoint(func, self.point_type(), **kwargs))
 
-    def create_time_series(self, label_values, func):
+    def create_time_series(self, label_values, func, **kwargs):
         """Create a derived measurement to trac `func`.
 
         :type label_values: list(:class:`LabelValue`)
@@ -432,7 +434,7 @@ def create_time_series(self, label_values, func):
             raise ValueError
         if func is None:
             raise ValueError
-        return self._create_time_series(label_values, func)
+        return self._create_time_series(label_values, func, **kwargs)
 
     def create_default_time_series(self, func):
         """Create the default derived measurement for this gauge.
diff --git a/tests/unit/metrics/export/test_gauge.py b/tests/unit/metrics/export/test_gauge.py
index 59fafd5ef..bd7ad791c 100644
--- a/tests/unit/metrics/export/test_gauge.py
+++ b/tests/unit/metrics/export/test_gauge.py
@@ -393,6 +393,15 @@ def test_create_time_series(self):
         unused_mock_fn.assert_not_called()
         self.assertEqual(len(derived_gauge.points.keys()), 1)
 
+        # with kwargs
+        def fn_with_args(value=None):
+            if value:
+                return value
+            return 0
+        label_values2 = [1, 2]
+        point3 = derived_gauge.create_time_series(label_values2, fn_with_args, value=5)  # noqa: E501
+        self.assertEqual(point3.get_value(), 5)
+
     def test_create_default_time_series(self):
         derived_gauge = gauge.DerivedLongGauge(
             Mock(), Mock(), Mock(), [Mock(), Mock])