Add request/response size sums

apitally · Jan 24, 2024 · a21fc2f · a21fc2f
1 parent 3d344d8
commit a21fc2f
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
diff --git a/apitally/client/base.py b/apitally/client/base.py
@@ -145,6 +145,8 @@ class RequestInfo:
 class RequestCounter:
     def __init__(self) -> None:
         self.request_counts: Counter[RequestInfo] = Counter()
+        self.request_size_sums: Counter[RequestInfo] = Counter()
+        self.response_size_sums: Counter[RequestInfo] = Counter()
         self.response_times: Dict[RequestInfo, Counter[int]] = {}
         self.request_sizes: Dict[RequestInfo, Counter[int]] = {}
         self.response_sizes: Dict[RequestInfo, Counter[int]] = {}
@@ -171,10 +173,14 @@ def add_request(
             self.request_counts[request_info] += 1
             self.response_times.setdefault(request_info, Counter())[response_time_ms_bin] += 1
             if request_size is not None:
-                request_size_kb_bin = int(request_size) // 1000  # In KB, rounded down to nearest 1KB
+                request_size = int(request_size)
+                request_size_kb_bin = request_size // 1000  # In KB, rounded down to nearest 1KB
+                self.request_size_sums[request_info] += request_size
                 self.request_sizes.setdefault(request_info, Counter())[request_size_kb_bin] += 1
             if response_size is not None:
-                response_size_kb_bin = int(response_size) // 1000  # In KB, rounded down to nearest 1KB
+                response_size = int(response_size)
+                response_size_kb_bin = response_size // 1000  # In KB, rounded down to nearest 1KB
+                self.response_size_sums[request_info] += response_size
                 self.response_sizes.setdefault(request_info, Counter())[response_size_kb_bin] += 1
 
     def get_and_reset_requests(self) -> List[Dict[str, Any]]:
@@ -188,12 +194,16 @@ def get_and_reset_requests(self) -> List[Dict[str, Any]]:
                         "path": request_info.path,
                         "status_code": request_info.status_code,
                         "request_count": count,
+                        "request_size_sum": self.request_size_sums.get(request_info, 0),
+                        "response_size_sum": self.response_size_sums.get(request_info, 0),
                         "response_times": self.response_times.get(request_info) or Counter(),
                         "request_sizes": self.request_sizes.get(request_info) or Counter(),
                         "response_sizes": self.response_sizes.get(request_info) or Counter(),
                     }
                 )
             self.request_counts.clear()
+            self.request_size_sums.clear()
+            self.response_size_sums.clear()
             self.response_times.clear()
             self.request_sizes.clear()
             self.response_sizes.clear()

diff --git a/tests/test_client_base.py b/tests/test_client_base.py
@@ -43,11 +43,15 @@ def test_request_counter():
     assert data[0]["path"] == "/test"
     assert data[0]["status_code"] == 200
     assert data[0]["request_count"] == 2
+    assert data[0]["request_size_sum"] == 0
+    assert data[0]["response_size_sum"] > 0
     assert data[0]["response_times"][100] == 1
     assert data[0]["response_times"][220] == 1
     assert len(data[0]["request_sizes"]) == 0
     assert data[0]["response_sizes"][0] == 2
     assert data[1]["method"] == "POST"
+    assert data[1]["request_size_sum"] > 0
+    assert data[1]["response_size_sum"] == 0
     assert data[1]["request_sizes"][0] == 1