From f42712e26477e366ff6407097c511a0bc88c9733 Mon Sep 17 00:00:00 2001
From: Arthur Wang <arthur@datadoghq.com>
Date: Tue, 9 Dec 2014 16:40:04 +0000
Subject: [PATCH] Add downtime update method and downtime message option

---
 src/dogapi/http/monitors.py            | 21 ++++++++++++++++++++-
 tests/integration/test_dog_http_api.py | 10 +++++++++-
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/src/dogapi/http/monitors.py b/src/dogapi/http/monitors.py
index 03be7d0..868e9fb 100644
--- a/src/dogapi/http/monitors.py
+++ b/src/dogapi/http/monitors.py
@@ -176,7 +176,7 @@ def unmute_monitor(self, monitor_id, scope=None):
 
 class DowntimeApi(object):
 
-    def schedule_downtime(self, scope, start, end=None):
+    def schedule_downtime(self, scope, start, end=None, message=None):
         """
         Schedule downtime over *scope* from *start* to *end*, where *start* and
         *end* are POSIX timestamps. If *end* is omitted then the downtime will
@@ -188,10 +188,29 @@ def schedule_downtime(self, scope, start, end=None):
         }
         if end:
             body['end'] = end
+        if message:
+            body['message'] = message
         return self.http_request('POST', '/downtime', body,
             response_formatter=lambda x: x['id'],
         )
 
+    def update_downtime(self, downtime_id, scope=None, start=None, end=None, message=None):
+        """
+        Update downtime parameters.
+        """
+        body = {}
+        if scope:
+            body['scope'] = scope
+        if start:
+            body['start'] = start
+        if end:
+            body['end'] = end
+        if message:
+            body['message'] = message
+        return self.http_request('PUT', '/downtime/%s' % downtime_id, body,
+            response_formatter=lambda x: x['id'],
+        )
+
     def get_downtime(self, downtime_id):
         """
         Get the downtime identified by *downtime_id*
diff --git a/tests/integration/test_dog_http_api.py b/tests/integration/test_dog_http_api.py
index 1189333..4286a05 100644
--- a/tests/integration/test_dog_http_api.py
+++ b/tests/integration/test_dog_http_api.py
@@ -688,11 +688,19 @@ def test_monitor_muting(self):
     def test_downtime(self):
         start = int(time.time())
         end = start + 1000
-        downtime_id = dog.schedule_downtime('env:staging', start, end)
+        downtime_id = dog.schedule_downtime('env:staging', start, end, "Message!")
         dt = dog.get_downtime(downtime_id)
         nt.assert_equal(dt['start'], start)
         nt.assert_equal(dt['end'], end)
         nt.assert_equal(dt['scope'], ['env:staging'])
+        nt.assert_equal(dt['message'], "Message!")
+
+        dog.update_downtime('env:prod', start + 1, end + 1, 'New Message!')
+        dt = dog.get_downtime(downtime_id)
+        nt.assert_equal(dt['start'], start + 1)
+        nt.assert_equal(dt['end'], end + 1)
+        nt.assert_equal(dt['scope'], ['env:prod'])
+        nt.assert_equal(dt['message'], "New Message!")
 
         dog.cancel_downtime(downtime_id)