From 0b665d3755f8c0467727786e282aab6462f37731 Mon Sep 17 00:00:00 2001
From: Rudi Giesler <rudi@praekelt.com>
Date: Wed, 4 Mar 2015 12:39:15 +0200
Subject: [PATCH] Added tests for outbound credit cutoff messages

---
 go/vumitools/billing_worker.py            | 11 ++++---
 go/vumitools/tests/test_billing_worker.py | 40 ++++++++++++++++++++++-
 2 files changed, 45 insertions(+), 6 deletions(-)

diff --git a/go/vumitools/billing_worker.py b/go/vumitools/billing_worker.py
index 076cb90f5..8bd9b6885 100644
--- a/go/vumitools/billing_worker.py
+++ b/go/vumitools/billing_worker.py
@@ -264,7 +264,7 @@ def process_outbound(self, config, msg, connector_name):
                 transaction = yield self.create_transaction_for_outbound(msg)
                 msg_mdh.set_paid()
                 if transaction.get('credit_cutoff_reached', False):
-                    self._handle_credit_cutoff(msg)
+                    msg = self._handle_credit_cutoff(msg)
         except BillingError:
             log.warning(
                 "BillingError for outbound message, sending without billing:"
@@ -280,11 +280,12 @@ def process_outbound(self, config, msg, connector_name):
                 msg, self.receive_inbound_connector, None)
 
     def _handle_credit_cutoff(self, msg):
-        if msg.session_event is not None:
-            msg.session_event = SESSION_CLOSE
-            msg.content = self.credit_limit_message
+        if msg.get('session_event') is not None:
+            msg['session_event'] = SESSION_CLOSE
+            msg['content'] = self.credit_limit_message
+            return msg
         else:
-            msg = None
+            return None
 
     @inlineCallbacks
     def process_event(self, config, event, connector_name):
diff --git a/go/vumitools/tests/test_billing_worker.py b/go/vumitools/tests/test_billing_worker.py
index 0bc12cb30..35187525f 100644
--- a/go/vumitools/tests/test_billing_worker.py
+++ b/go/vumitools/tests/test_billing_worker.py
@@ -5,6 +5,7 @@
 from twisted.internet.defer import inlineCallbacks, returnValue
 from twisted.web.client import Agent, Request, Response
 
+from vumi.message import TransportUserMessage
 from vumi.tests.helpers import VumiTestCase
 from vumi.tests.utils import LogCatcher
 from vumi.utils import mkheaders, StringProducer
@@ -17,11 +18,14 @@
 from go.billing.api import BillingError
 from go.billing.utils import JSONEncoder
 
+SESSION_CLOSE = TransportUserMessage.SESSION_CLOSE
+
 
 class BillingApiMock(object):
 
-    def __init__(self):
+    def __init__(self, credit_cutoff=False):
         self.transactions = []
+        self.credit_cutoff = credit_cutoff
 
     def _record(self, items, vars):
         del vars["self"]
@@ -50,6 +54,7 @@ def create_transaction(self, account_number, message_id, tag_pool_name,
             "status": "Completed",
             "transaction_type": transaction_type,
             "session_length": session_length,
+            "credit_cutoff_reached": self.credit_cutoff
         }
 
 
@@ -739,3 +744,36 @@ def test_outbound_message_session_length_custom_field(self):
             "outbound",
             session_created=False,
             session_metadata_field='foo')
+
+    @inlineCallbacks
+    def test_outbound_message_credit_cutoff_session(self):
+        self.billing_api = BillingApiMock(credit_cutoff=True)
+        dispatcher = yield self.get_dispatcher()
+
+        yield self.make_dispatch_outbound(
+            "outbound",
+            user_account="12345",
+            tag=("pool1", "1234"),
+            helper_metadata={},
+            session_event='new'
+            )
+
+        published_msg = self.ri_helper.get_dispatched_outbound()[0]
+        self.assertEqual(published_msg['session_event'], SESSION_CLOSE)
+        self.assertEqual(
+            published_msg['content'], dispatcher.credit_limit_message)
+
+    @inlineCallbacks
+    def test_outbound_message_credit_cutoff_message(self):
+        self.billing_api = BillingApiMock(credit_cutoff=True)
+        yield self.get_dispatcher()
+
+        yield self.make_dispatch_outbound(
+            "outbound",
+            user_account="12345",
+            tag=("pool1", "1234"),
+            helper_metadata={},
+            )
+
+        published_msgs = self.ri_helper.get_dispatched_outbound()
+        self.assertEqual(len(published_msgs), 0)