feat: add quick stats option

TOD-theses · Aug 6, 2024 · 02c7159 · 02c7159
1 parent 85ff77d
commit 02c7159
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 8 deletions.
diff --git a/tests/integration/test_snapshots.py b/tests/integration/test_snapshots.py
@@ -1,4 +1,3 @@
-from pathlib import Path
 
 import pytest
 from tod_attack_miner.db.db import DB
@@ -15,7 +14,6 @@
 from tod_attack_miner.rpc.rpc import RPC
 
 test_provider_url = "http://localhost:8124/eth"
-test_db_path = Path("test_database.db")
 
 
 @pytest.mark.vcr

diff --git a/tod_attack_miner/cli.py b/tod_attack_miner/cli.py
@@ -39,6 +39,9 @@ def main():
     parser.add_argument("--postgres-password", type=str, default="password")
     parser.add_argument("--postgres-host", type=str, default="localhost")
     parser.add_argument("--postgres-port", type=int, default=5432)
+    parser.add_argument(
+        "--quick-stats", action="store_true", help="Only output performant stats"
+    )
     parser.add_argument(
         "--stats-only",
         action=BooleanOptionalAction,
@@ -52,7 +55,7 @@ def main():
         miner = Miner(RPC(args.archive_node_provider), DB(conn))
 
         if args.stats_only:
-            print(json.dumps(miner.get_stats()))
+            print(json.dumps(miner.get_stats(args.quick_stats)))
         else:
             if args.reset_db:
                 miner.reset_db()
@@ -64,4 +67,4 @@ def main():
                 + get_filters_duplicate_limits(10)
             )
             print(f"Found {miner.count_candidates()} candidates")
-            print(json.dumps(miner.get_stats()))
+            print(json.dumps(miner.get_stats(args.quick_stats)))
diff --git a/tod_attack_miner/miner/miner.py b/tod_attack_miner/miner/miner.py
@@ -39,10 +39,15 @@ def count_candidates(self) -> int:
     def get_candidates(self) -> Sequence[Candidate]:
         return self.db.get_candidates()
 
-    def get_stats(self):
-        self._filter_stats["candidates"]["original_without_same_value"] = (
-            self.db.count_candidates_original()
-        )
+    def get_stats(self, quick=False):
+        if quick:
+            self._filter_stats["candidates"]["original_without_same_value"] = (
+                "<omitted because of quick stats>"
+            )
+        else:
+            self._filter_stats["candidates"]["original_without_same_value"] = (
+                self.db.count_candidates_original()
+            )
         return {
             "accesses": self.db.get_accesses_stats(),
             "state_diffs": self.db.get_state_diffs_stats(),