Methods for getting the unique address counts.

vumi_message_store/batch_info_cache.py

@@ -335,6 +335,18 @@ def get_event_count(self, batch_id):
         """
         return self._get_counter_value(self.event_count_key(batch_id))
 
+    def get_from_addr_count(self, batch_id):
+        """
+        Return the count of from addresses.
+        """
+        return self.redis.pfcount(self.from_addr_key(batch_id))
+
+    def get_to_addr_count(self, batch_id):
+        """
+        Return the count of to addresses.
+        """
+        return self.redis.pfcount(self.to_addr_key(batch_id))
+
     @Manager.calls_manager
     def rebuild_cache(self, batch_id, qms, page_size=None):
         """

vumi_message_store/message_store.py

@@ -263,3 +263,9 @@ def get_batch_outbound_count(self, batch_id):
 
     def get_batch_event_count(self, batch_id):
         return self.batch_info_cache.get_event_count(batch_id)
+
+    def get_batch_from_addr_count(self, batch_id):
+        return self.batch_info_cache.get_from_addr_count(batch_id)
+
+    def get_batch_to_addr_count(self, batch_id):
+        return self.batch_info_cache.get_to_addr_count(batch_id)

vumi_message_store/tests/test_batch_info_cache.py

@@ -277,7 +277,7 @@ def test_add_from_addrs(self):
         Adding a from_addr updates the HyperLogLog counter for the batch.
         """
         yield self.batch_info_cache.batch_start("mybatch")
-        incr = yield self.batch_info_cache.add_from_addr("mybatch", "from-1")
+        incr = yield self.batch_info_cache.add_from_addr("mybatch", "addr-1")
         self.assertEqual(incr, 1)
 
         yield self.assert_redis_keys([
@@ -291,12 +291,12 @@ def test_add_from_addrs(self):
         yield self.assert_redis_pfcount("batches:from_addr_hll:mybatch", 1)
 
         # Adding a second address updates the counter.
-        incr = yield self.batch_info_cache.add_from_addr("mybatch", "from-2")
+        incr = yield self.batch_info_cache.add_from_addr("mybatch", "addr-2")
         self.assertEqual(incr, 1)
         yield self.assert_redis_pfcount("batches:from_addr_hll:mybatch", 2)
 
         # Adding a previously-added address doesn't update the counter.
-        incr = yield self.batch_info_cache.add_from_addr("mybatch", "from-1")
+        incr = yield self.batch_info_cache.add_from_addr("mybatch", "addr-1")
         self.assertEqual(incr, 0)
         yield self.assert_redis_pfcount("batches:from_addr_hll:mybatch", 2)
 
@@ -432,7 +432,7 @@ def test_add_to_addrs(self):
         Adding a to_addr updates the HyperLogLog counter for the batch.
         """
         yield self.batch_info_cache.batch_start("mybatch")
-        incr = yield self.batch_info_cache.add_to_addr("mybatch", "from-1")
+        incr = yield self.batch_info_cache.add_to_addr("mybatch", "addr-1")
         self.assertEqual(incr, 1)
 
         yield self.assert_redis_keys([
@@ -446,12 +446,12 @@ def test_add_to_addrs(self):
         yield self.assert_redis_pfcount("batches:to_addr_hll:mybatch", 1)
 
         # Adding a second address updates the counter.
-        incr = yield self.batch_info_cache.add_to_addr("mybatch", "from-2")
+        incr = yield self.batch_info_cache.add_to_addr("mybatch", "addr-2")
         self.assertEqual(incr, 1)
         yield self.assert_redis_pfcount("batches:to_addr_hll:mybatch", 2)
 
         # Adding a previously-added address doesn't update the counter.
-        incr = yield self.batch_info_cache.add_to_addr("mybatch", "from-1")
+        incr = yield self.batch_info_cache.add_to_addr("mybatch", "addr-1")
         self.assertEqual(incr, 0)
         yield self.assert_redis_pfcount("batches:to_addr_hll:mybatch", 2)
 
@@ -986,6 +986,48 @@ def test_get_event_count_no_batch(self):
         count = yield self.batch_info_cache.get_event_count("batch")
         self.assertEqual(count, 0)
 
+    @inlineCallbacks
+    def test_get_from_addr_count(self):
+        """
+        The from_addr count can be queried.
+        """
+        yield self.batch_info_cache.batch_start("batch")
+        yield self.batch_info_cache.add_from_addr("batch", "addr-1")
+        yield self.batch_info_cache.add_from_addr("batch", "addr-2")
+        yield self.batch_info_cache.add_from_addr("batch", "addr-3")
+
+        count = yield self.batch_info_cache.get_from_addr_count("batch")
+        self.assertEqual(count, 3)
+
+    @inlineCallbacks
+    def test_get_from_addr_count_no_batch(self):
+        """
+        The from_addr count returns zero for missing batches.
+        """
+        count = yield self.batch_info_cache.get_from_addr_count("batch")
+        self.assertEqual(count, 0)
+
+    @inlineCallbacks
+    def test_get_to_addr_count(self):
+        """
+        The to_addr count can be queried.
+        """
+        yield self.batch_info_cache.batch_start("batch")
+        yield self.batch_info_cache.add_to_addr("batch", "addr-1")
+        yield self.batch_info_cache.add_to_addr("batch", "addr-2")
+        yield self.batch_info_cache.add_to_addr("batch", "addr-3")
+
+        count = yield self.batch_info_cache.get_to_addr_count("batch")
+        self.assertEqual(count, 3)
+
+    @inlineCallbacks
+    def test_get_to_addr_count_no_batch(self):
+        """
+        The to_addr count returns zero for missing batches.
+        """
+        count = yield self.batch_info_cache.get_to_addr_count("batch")
+        self.assertEqual(count, 0)
+
     @inlineCallbacks
     def test_rebuild_cache(self):
         """

vumi_message_store/tests/test_message_store.py

@@ -1464,3 +1464,45 @@ def test_get_batch_event_count_no_batch(self):
         """
         count = yield self.store.get_batch_event_count("batch")
         self.assertEqual(count, 0)
+
+    @inlineCallbacks
+    def test_get_batch_from_addr_count(self):
+        """
+        The from_addr count can be queried.
+        """
+        yield self.bi_cache.batch_start("batch")
+        yield self.bi_cache.add_from_addr("batch", "addr-1")
+        yield self.bi_cache.add_from_addr("batch", "addr-2")
+        yield self.bi_cache.add_from_addr("batch", "addr-3")
+
+        count = yield self.store.get_batch_from_addr_count("batch")
+        self.assertEqual(count, 3)
+
+    @inlineCallbacks
+    def test_get_batch_from_addr_count_no_batch(self):
+        """
+        The from_addr count returns zero for missing batches.
+        """
+        count = yield self.store.get_batch_from_addr_count("batch")
+        self.assertEqual(count, 0)
+
+    @inlineCallbacks
+    def test_get_batch_to_addr_count(self):
+        """
+        The to_addr count can be queried.
+        """
+        yield self.bi_cache.batch_start("batch")
+        yield self.bi_cache.add_to_addr("batch", "addr-1")
+        yield self.bi_cache.add_to_addr("batch", "addr-2")
+        yield self.bi_cache.add_to_addr("batch", "addr-3")
+
+        count = yield self.store.get_batch_to_addr_count("batch")
+        self.assertEqual(count, 3)
+
+    @inlineCallbacks
+    def test_get_batch_to_addr_count_no_batch(self):
+        """
+        The to_addr count returns zero for missing batches.
+        """
+        count = yield self.store.get_batch_to_addr_count("batch")
+        self.assertEqual(count, 0)