From 820eb6878983664e707041be8e152b0c4ef1dd9f Mon Sep 17 00:00:00 2001
From: mmsqe <mavis@crypto.com>
Date: Fri, 13 Sep 2024 11:23:50 +0800
Subject: [PATCH] Apply suggestions from code review

---
 testground/benchmark/benchmark/stateless.py | 45 ++++++++++++---------
 1 file changed, 27 insertions(+), 18 deletions(-)

diff --git a/testground/benchmark/benchmark/stateless.py b/testground/benchmark/benchmark/stateless.py
index 624b9459ad..b7c6493ac6 100644
--- a/testground/benchmark/benchmark/stateless.py
+++ b/testground/benchmark/benchmark/stateless.py
@@ -321,21 +321,29 @@ def truncate_fractional_seconds(timestamp):
     return timestamp
 
 
-def calculate_tps(height, num_txs, timestamp):
-    if height > 1:
-        prev_block = block(height - 1)
-        prev_timestamp = prev_block["result"]["block"]["header"]["time"]
-        time_format = "%Y-%m-%dT%H:%M:%S.%fZ"
-        current_time, previous_time = (
-            datetime.strptime(truncate_fractional_seconds(ts), time_format)
-            for ts in (timestamp, prev_timestamp)
-        )
-        diff = (current_time - previous_time).total_seconds()
-        if diff > 0:
-            tps = num_txs / diff
-            return tps
-        else:
-            return 0
+def calculate_tps(start_height, end_height):
+    total_txs = 0
+    total_time = 0
+
+    for height in range(start_height + 1, end_height):
+        blk = block(height)
+        timestamp = blk["result"]["block"]["header"]["time"]
+        prev_blk = block(height - 1)
+        txs = len(prev_blk["result"]["block"]["data"]["txs"])
+        if txs > 0:
+            prev_timestamp = prev_blk["result"]["block"]["header"]["time"]
+            time_format = "%Y-%m-%dT%H:%M:%S.%fZ"
+            current_time, previous_time = (
+                datetime.strptime(truncate_fractional_seconds(ts), time_format)
+                for ts in (timestamp, prev_timestamp)
+            )
+            diff = (current_time - previous_time).total_seconds()
+            total_time += diff
+            total_txs += txs
+
+    if total_time > 0:
+        average_tps = total_txs / total_time
+        return average_tps
     else:
         return 0
 
@@ -345,13 +353,14 @@ def dump_block_stats(fp):
     dump simple statistics for blocks for analysis
     """
     tps_list = []
-    for i in range(1, block_height() + 1):
+    current = block_height()
+    for i in range(1, current + 1):
         blk = block(i)
         timestamp = blk["result"]["block"]["header"]["time"]
         txs = len(blk["result"]["block"]["data"]["txs"])
-        tps = calculate_tps(i, txs, timestamp)
-        print("block", i, txs, timestamp, tps, file=fp)
+        tps = calculate_tps(i, min(i + 2, current))
         tps_list.append(tps)
+        print("block", i, txs, timestamp, tps, file=fp)
     top_tps = sorted(tps_list, reverse=True)[:3]
     print("top_tps", top_tps)