From a96d246e11a046bd58fe44392a3cce7522c741f4 Mon Sep 17 00:00:00 2001
From: Florian Pellissier <111426680+flopell@users.noreply.github.com>
Date: Fri, 14 Jun 2024 17:33:29 +0200
Subject: [PATCH] fix(test runner): fix how we check balances and stout full
 buffer

---
 testing/runner.py | 22 +++++++++++++++-------
 testing/tycho.py  | 20 ++++++++++----------
 2 files changed, 25 insertions(+), 17 deletions(-)

diff --git a/testing/runner.py b/testing/runner.py
index 195c322..a9136ed 100644
--- a/testing/runner.py
+++ b/testing/runner.py
@@ -98,14 +98,22 @@ class TestRunner:
                         return TestResult.Failed(
                             f"Value mismatch for key '{key}': {value} != {component[key]}"
                         )
-                for state in protocol_states["states"]:
-                    for token, balance in state["balances"].items():
-                        node_balance = get_token_balance(token,comp_id,stop_block)
-                        tycho_balance = int(balance,16)
-                        if node_balance != tycho_balance:
-                            return TestResult.Failed(f"Balance mismatch for {comp_id}:{token} at block {stop_block}: got {node_balance} from rpc call and {tycho_balance} from Substreams")
-            return TestResult.Passed()
 
+            for component in protocol_components["protocol_components"]:
+                comp_id = component["id"].lower()
+                for token in component["tokens"]:
+                    token_lower = token.lower()
+                    state = next((s for s in protocol_states["states"] if s["component_id"].lower() == comp_id), None)
+                    if state:
+                        balance_hex = state["balances"].get(token_lower, "0x0")
+                    else:
+                        balance_hex = "0x0"
+
+                    node_balance = get_token_balance(token, comp_id, stop_block)
+                    tycho_balance = int(balance_hex, 16)
+                    if node_balance != tycho_balance:
+                        return TestResult.Failed(f"Balance mismatch for {comp_id}:{token} at block {stop_block}: got {node_balance} from rpc call and {tycho_balance} from Substreams")
+            return TestResult.Passed()
         except Exception as e:
             return TestResult.Failed(str(e))
 
diff --git a/testing/tycho.py b/testing/tycho.py
index 1806b0a..8408cc7 100644
--- a/testing/tycho.py
+++ b/testing/tycho.py
@@ -41,23 +41,23 @@ class TychoRunner:
                     str(start_block),
                     "--stop-block",
                     str(end_block + 2),
-                ],  # TODO: +2 is a hack to make up for the cache in the index side.
+                ],  # +2 is to make up for the cache in the index side.
                 stdout=subprocess.PIPE,
                 stderr=subprocess.PIPE,
                 text=True,
                 bufsize=1,
                 env=env,
             )
-            if self.with_binary_logs:
-                with process.stdout:
-                    for line in iter(process.stdout.readline, ""):
-                        if line:
-                            print(line.strip())
+            
+            with process.stdout:
+                for line in iter(process.stdout.readline, ""):
+                    if line and self.with_binary_logs:
+                        print(line.strip())
 
-                with process.stderr:
-                    for line in iter(process.stderr.readline, ""):
-                        if line:
-                            print(line.strip())
+            with process.stderr:
+                for line in iter(process.stderr.readline, ""):
+                    if line and self.with_binary_logs:
+                        print(line.strip())
 
             process.wait()