From ce523d79273d7b70e667f179a8bba9656758fcb0 Mon Sep 17 00:00:00 2001
From: Jason Derrett <jason@librato.com>
Date: Tue, 5 Jul 2016 10:08:57 -0500
Subject: [PATCH] Add gap_detection attribute to Stream and allow extra
 attributes

---
 librato/streams.py    |  9 +++++++--
 tests/test_streams.py | 18 ++++++++++++++++++
 2 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/librato/streams.py b/librato/streams.py
index 0f2fdc2..5cdf2a1 100644
--- a/librato/streams.py
+++ b/librato/streams.py
@@ -3,11 +3,11 @@ def __init__(self, metric=None, source='*', composite=None,
                  name=None, type=None, id=None,
                  group_function=None, summary_function=None,
                  transform_function=None, downsample_function=None,
-                 period=None, split_axis=None,
+                 period=None, split_axis=None, gap_detection=None,
                  min=None, max=None,
                  units_short=None, units_long=None, color=None,
                  # deprecated
-                 composite_function=None
+                 composite_function=None, **kwargs
                  ):
         self.metric = metric
         self.source = source
@@ -30,6 +30,11 @@ def __init__(self, metric=None, source='*', composite=None,
         self.units_short = units_short
         self.units_long = units_long
         self.color = color
+        self.gap_detection = gap_detection
+
+        # Pick up any attributes that are not explicitly defined
+        for attr in kwargs:
+            setattr(self, attr, kwargs[attr])
 
         # Can't have a composite and source/metric
         if self.composite:
diff --git a/tests/test_streams.py b/tests/test_streams.py
index a9c494e..7dd783d 100644
--- a/tests/test_streams.py
+++ b/tests/test_streams.py
@@ -74,6 +74,24 @@ def test_init_units(self):
         self.assertEqual(Stream(units_short='req/s').units_short, 'req/s')
         self.assertEqual(Stream(units_long='requests per second').units_long, 'requests per second')
 
+    def test_init_color(self):
+        self.assertIsNone(Stream().color)
+        self.assertEqual(Stream(color='#f00').color, '#f00')
+
+    def test_init_gap_detection(self):
+        self.assertIsNone(Stream().gap_detection)
+        self.assertTrue(Stream(gap_detection=True).gap_detection)
+        self.assertFalse(Stream(gap_detection=False).gap_detection)
+
+    # Adding this to avoid exceptions raised due to unknown Stream attributes
+    def test_init_with_extra_attributes(self):
+        attrs = {"color": "#f00", "something": "foo"}
+        s = Stream(**attrs)
+        # color is a known attribute
+        self.assertEqual(s.color, '#f00')
+        self.assertEqual(s.something, 'foo')
+
+
     def test_get_payload(self):
         self.assertEqual(Stream(metric='my.metric').get_payload(),
             {'metric': 'my.metric', 'source': '*'})