Add flow_count and active_flow_count methods to flow.State.

Use these in unit tests where sensible.

1 files changed, 19 insertions, 9 deletions

diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 3c62169f..857c8ae4 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -164,6 +164,15 @@ class Flow:
         self.intercepting = False
         self._backup = None
 
+    @classmethod
+    def from_state(klass, state):
+        f = klass(None)
+        f.load_state(state)
+        return f
+
+    def __eq__(self, other):
+        return self.get_state() == other.get_state()
+
     def script_serialize(self):
         data = self.get_state()
         return json.dumps(data)
@@ -246,15 +255,6 @@ class Flow:
         else:
             self.error = None
 
-    @classmethod
-    def from_state(klass, state):
-        f = klass(None)
-        f.load_state(state)
-        return f
-
-    def __eq__(self, other):
-        return self.get_state() == other.get_state()
-
     def modified(self):
         # FIXME: Save a serialization in backup, compare current with
         # backup to detect if flow has _really_ been modified.
@@ -308,6 +308,16 @@ class State:
         self.limit = None
         self.intercept = None
 
+    def flow_count(self):
+        return len(self.flow_map)
+
+    def active_flow_count(self):
+        c = 0
+        for i in self.flow_list:
+            if not i.response and not i.error:
+                c += 1
+        return c
+
     def clientconnect(self, cc):
         self.client_connections.append(cc)