From 18efbfb72d5bd2f32e6f115b24ce9da159a2a789 Mon Sep 17 00:00:00 2001
From: Lutz Ostkamp <lutz.ostkamp@flixbus.com>
Date: Thu, 2 Dec 2021 15:34:18 +0100
Subject: [PATCH] Fix failing test from previos PR

---
 airflow/stats.py         | 9 ++++++---
 tests/core/test_stats.py | 5 +++++
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/airflow/stats.py b/airflow/stats.py
index e3b069bd99ce1..35bf84409ce15 100644
--- a/airflow/stats.py
+++ b/airflow/stats.py
@@ -16,12 +16,13 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import datetime
 import logging
 import socket
 import string
 import time
 from functools import wraps
-from typing import TYPE_CHECKING, Callable, Optional, TypeVar, cast
+from typing import TYPE_CHECKING, Callable, List, Optional, TypeVar, Union, cast
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowConfigException, InvalidStatsNameException
@@ -64,7 +65,7 @@ def gauge(cls, stat: str, value: float, rate: int = 1, delta: bool = False) -> N
         """Gauge stat"""
 
     @classmethod
-    def timing(cls, stat: str, dt) -> None:
+    def timing(cls, stat: str, dt: Union[float, datetime.timedelta]) -> None:
         """Stats timing"""
 
     @classmethod
@@ -317,10 +318,12 @@ def gauge(self, stat, value, rate=1, delta=False, tags=None):
         return None
 
     @validate_stat
-    def timing(self, stat, dt, tags=None):
+    def timing(self, stat, dt: Union[float, datetime.timedelta], tags: Optional[List[str]] = None):
         """Stats timing"""
         if self.allow_list_validator.test(stat):
             tags = tags or []
+            if isinstance(dt, datetime.timedelta):
+                dt = dt.total_seconds()
             return self.dogstatsd.timing(metric=stat, value=dt, tags=tags)
         return None
 
diff --git a/tests/core/test_stats.py b/tests/core/test_stats.py
index 83169e2935b24..c401a2f61ef95 100644
--- a/tests/core/test_stats.py
+++ b/tests/core/test_stats.py
@@ -181,9 +181,14 @@ def test_empty_timer(self):
         self.dogstatsd_client.timed.assert_not_called()
 
     def test_timing(self):
+        import datetime
+
         self.dogstatsd.timing("dummy_timer", 123)
         self.dogstatsd_client.timing.assert_called_once_with(metric='dummy_timer', value=123, tags=[])
 
+        self.dogstatsd.timing("dummy_timer", datetime.timedelta(seconds=123))
+        self.dogstatsd_client.timing.assert_called_with(metric='dummy_timer', value=123.0, tags=[])
+
     def test_gauge(self):
         self.dogstatsd.gauge("dummy", 123)
         self.dogstatsd_client.gauge.assert_called_once_with(metric='dummy', sample_rate=1, value=123, tags=[])