From 9a6e8c5c7c1e7d4077d5e27610a26d30147e109e Mon Sep 17 00:00:00 2001
From: "google-labs-jules[bot]"
 <161369871+google-labs-jules[bot]@users.noreply.github.com>
Date: Sat, 21 Mar 2026 08:23:16 +0000
Subject: [PATCH] Add edge case tests for _find_col in ttm_analysis

Co-authored-by: aguzererler <6199053+aguzererler@users.noreply.github.com>
---
 tests/unit/test_ttm_analysis.py | 20 +++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

diff --git a/tests/unit/test_ttm_analysis.py b/tests/unit/test_ttm_analysis.py
index 342b1928..2cdede41 100644
--- a/tests/unit/test_ttm_analysis.py
+++ b/tests/unit/test_ttm_analysis.py
@@ -2,7 +2,6 @@
 
 import pytest
 import pandas as pd
-from io import StringIO
 
 
 # ---------------------------------------------------------------------------
@@ -52,6 +51,25 @@ def _make_cashflow_csv(n_quarters: int = 8) -> str:
     return df.to_csv()
 
 
+# ---------------------------------------------------------------------------
+# Unit tests for _find_col
+# ---------------------------------------------------------------------------
+
+class TestFindCol:
+    def setup_method(self):
+        from tradingagents.dataflows.ttm_analysis import _find_col
+        self.find_col = _find_col
+
+    def test_find_col_success(self):
+        df = pd.DataFrame({"Total Revenue": [100], "Gross Profit": [40]})
+        candidates = ["Revenue", "Total Revenue", "totalRevenue"]
+        assert self.find_col(df, candidates) == "Total Revenue"
+
+    def test_find_col_not_found(self):
+        df = pd.DataFrame({"Unrelated": [100], "Another": [40]})
+        candidates = ["Revenue", "Total Revenue", "totalRevenue"]
+        assert self.find_col(df, candidates) is None
+
 # ---------------------------------------------------------------------------
 # Unit tests for compute_ttm_metrics
 # ---------------------------------------------------------------------------