[BugFix] Fix Equity Integration Tests (#6086)

* fix equity integration tests * black * pylint * black * black * move df_apply to helpers * unused import * remove file that shouldn't have been added * oops. didn't remove a couple lines
2026-10-27 12:40:35 +00:00 · 2024-02-18 09:23:30 -08:00
parent ccf46aad92
commit 60a93c3266
10 changed files with 147 additions and 132 deletions
--- a/openbb_platform/extensions/equity/integration/test_equity_python.py
+++ b/openbb_platform/extensions/equity/integration/test_equity_python.py
@@ -1225,9 +1225,7 @@ def test_equity_price_quote(params, obb):
@parametrize(
    "params",
    [
-        ({"symbol": "AAPL", "provider": "cboe"}),
        ({"symbol": "MSFT", "provider": "intrinio"}),
-        ({"symbol": "AAPL,MSFT", "provider": "cboe"}),
        ({"symbol": "AAPL,MSFT", "provider": "intrinio"}),
        ({"symbol": "AAPL,MSFT", "provider": "finviz"}),
        ({"symbol": "AAPL,MSFT", "provider": "yfinance"}),
--- a/openbb_platform/providers/finviz/openbb_finviz/models/price_performance.py
+++ b/openbb_platform/providers/finviz/openbb_finviz/models/price_performance.py
@@ -20,6 +20,8 @@ class FinvizPricePerformanceQueryParams(RecentPerformanceQueryParams):
    Source: https://finviz.com/screener.ashx
    """

+    __json_schema_extra__ = {"symbol": ["multiple_items_allowed"]}
+

 class FinvizPricePerformanceData(RecentPerformanceData):
    """Finviz Price Performance Data."""
--- a/openbb_platform/providers/yfinance/openbb_yfinance/models/active.py
+++ b/openbb_platform/providers/yfinance/openbb_yfinance/models/active.py
@@ -1,15 +1,18 @@
 """Yahoo Finance Asset Performance Active Model."""

+# pylint: disable=unused-argument
+
 import re
 from typing import Any, Dict, List, Optional

-import pandas as pd
-import requests
 from openbb_core.provider.abstract.fetcher import Fetcher
 from openbb_core.provider.standard_models.equity_performance import (
    EquityPerformanceData,
    EquityPerformanceQueryParams,
 )
+from openbb_core.provider.utils.helpers import make_request
+from openbb_yfinance.utils.helpers import df_transform_numbers
+from pandas import DataFrame, read_html
 from pydantic import Field


@@ -35,10 +38,10 @@ class YFActiveData(EquityPerformanceData):
        "pe_ratio_ttm": "PE Ratio (TTM)",
    }

-    market_cap: str = Field(
+    market_cap: Optional[float] = Field(
        description="Market Cap displayed in billions.",
    )
-    avg_volume_3_months: float = Field(
+    avg_volume_3_months: Optional[float] = Field(
        description="Average volume over the last 3 months in millions.",
    )
    pe_ratio_ttm: Optional[float] = Field(
@@ -60,35 +63,31 @@ class YFActiveFetcher(Fetcher[YFActiveQueryParams, List[YFActiveData]]):
        query: YFActiveQueryParams,
        credentials: Optional[Dict[str, str]],
        **kwargs: Any,
-    ) -> pd.DataFrame:
+    ) -> DataFrame:
        """Get data from YF."""
        headers = {"user_agent": "Mozilla/5.0"}
-        html = requests.get(
+        html = make_request(
            "https://finance.yahoo.com/screener/predefined/most_actives",
            headers=headers,
-            timeout=10,
        ).text
        html_clean = re.sub(r"(<span class=\"Fz\(0\)\">).*?(</span>)", "", html)
-        df = (
-            pd.read_html(html_clean, header=None)[0]
-            .dropna(how="all", axis=1)
-            .fillna("-")
-            .replace("-", None)
-        )
+        df = read_html(html_clean, header=None)[0].dropna(how="all", axis=1)
        return df

    @staticmethod
    def transform_data(
        query: EquityPerformanceQueryParams,
-        data: List[Dict],
+        data: DataFrame,
        **kwargs: Any,
    ) -> List[YFActiveData]:
        """Transform data."""
-        data["% Change"] = data["% Change"].str.replace("%", "")
-        data["Volume"] = data["Volume"].str.replace("M", "").astype(float) * 1000000
-        data["Avg Vol (3 month)"] = (
-            data["Avg Vol (3 month)"].str.replace("M", "").astype(float) * 1000000
-        )
-        data = data.to_dict(orient="records")
-        data = sorted(data, key=lambda d: d["Volume"], reverse=query.sort == "desc")
-        return [YFActiveData.model_validate(d) for d in data]
+
+        columns = ["Market Cap", "Avg Vol (3 month)", "Volume", "% Change"]
+
+        data = df_transform_numbers(data, columns)
+        data = data.fillna("N/A").replace("N/A", None)
+
+        return [
+            YFActiveData.model_validate(d)
+            for d in data.sort_values(by="Volume", ascending=False).to_dict("records")
+        ]
--- a/openbb_platform/providers/yfinance/openbb_yfinance/models/aggressive_small_caps.py
+++ b/openbb_platform/providers/yfinance/openbb_yfinance/models/aggressive_small_caps.py
@@ -1,15 +1,18 @@
 """Yahoo Finance Aggressive Small Caps Model."""

+# pylint: disable=unused-argument
+
 import re
 from typing import Any, Dict, List, Optional

 import pandas as pd
-import requests
 from openbb_core.provider.abstract.fetcher import Fetcher
 from openbb_core.provider.standard_models.equity_performance import (
    EquityPerformanceData,
    EquityPerformanceQueryParams,
 )
+from openbb_core.provider.utils.helpers import make_request
+from openbb_yfinance.utils.helpers import df_transform_numbers
 from pydantic import Field


@@ -35,10 +38,12 @@ class YFAggressiveSmallCapsData(EquityPerformanceData):
        "pe_ratio_ttm": "PE Ratio (TTM)",
    }

-    market_cap: str = Field(
+    market_cap: Optional[float] = Field(
+        default=None,
        description="Market Cap.",
    )
-    avg_volume_3_months: float = Field(
+    avg_volume_3_months: Optional[float] = Field(
+        default=None,
        description="Average volume over the last 3 months in millions.",
    )
    pe_ratio_ttm: Optional[float] = Field(
@@ -65,18 +70,12 @@ class YFAggressiveSmallCapsFetcher(
    ) -> pd.DataFrame:
        """Get data from YF."""
        headers = {"user_agent": "Mozilla/5.0"}
-        html = requests.get(
+        html = make_request(
            "https://finance.yahoo.com/screener/predefined/aggressive_small_caps",
            headers=headers,
-            timeout=10,
        ).text
        html_clean = re.sub(r"(<span class=\"Fz\(0\)\">).*?(</span>)", "", html)
-        df = (
-            pd.read_html(html_clean, header=None)[0]
-            .dropna(how="all", axis=1)
-            .fillna("-")
-            .replace("-", None)
-        )
+        df = pd.read_html(html_clean, header=None)[0].dropna(how="all", axis=1)
        return df

    @staticmethod
@@ -86,11 +85,11 @@ class YFAggressiveSmallCapsFetcher(
        **kwargs: Any,
    ) -> List[YFAggressiveSmallCapsData]:
        """Transform data."""
-        data["% Change"] = data["% Change"].str.replace("%", "")
-        data["Volume"] = data["Volume"].str.replace("M", "").astype(float) * 1000000
-        data["Avg Vol (3 month)"] = (
-            data["Avg Vol (3 month)"].str.replace("M", "").astype(float) * 1000000
-        )
-        data = data.to_dict(orient="records")
-        data = sorted(data, key=lambda d: d["% Change"], reverse=query.sort == "desc")
-        return [YFAggressiveSmallCapsData.model_validate(d) for d in data]
+
+        columns = ["Market Cap", "Avg Vol (3 month)", "Volume", "% Change"]
+
+        data = df_transform_numbers(data=data, columns=columns)
+        data = data.fillna("N/A").replace("N/A", None)
+        return [
+            YFAggressiveSmallCapsData.model_validate(d) for d in data.to_dict("records")
+        ]
--- a/openbb_platform/providers/yfinance/openbb_yfinance/models/gainers.py
+++ b/openbb_platform/providers/yfinance/openbb_yfinance/models/gainers.py
@@ -1,15 +1,18 @@
 """Yahoo Finance Asset Performance Gainers Model."""

+# pylint: disable=unused-argument
+
 import re
 from typing import Any, Dict, List, Optional

 import pandas as pd
-import requests
 from openbb_core.provider.abstract.fetcher import Fetcher
 from openbb_core.provider.standard_models.equity_performance import (
    EquityPerformanceData,
    EquityPerformanceQueryParams,
 )
+from openbb_core.provider.utils.helpers import make_request
+from openbb_yfinance.utils.helpers import df_transform_numbers
 from pydantic import Field


@@ -35,10 +38,10 @@ class YFGainersData(EquityPerformanceData):
        "pe_ratio_ttm": "PE Ratio (TTM)",
    }

-    market_cap: str = Field(
+    market_cap: Optional[float] = Field(
        description="Market Cap.",
    )
-    avg_volume_3_months: float = Field(
+    avg_volume_3_months: Optional[float] = Field(
        description="Average volume over the last 3 months in millions.",
    )
    pe_ratio_ttm: Optional[float] = Field(
@@ -63,18 +66,12 @@ class YFGainersFetcher(Fetcher[YFGainersQueryParams, List[YFGainersData]]):
    ) -> pd.DataFrame:
        """Get data from YF."""
        headers = {"user_agent": "Mozilla/5.0"}
-        html = requests.get(
+        html = make_request(
            "https://finance.yahoo.com/screener/predefined/day_gainers",
            headers=headers,
-            timeout=10,
        ).text
        html_clean = re.sub(r"(<span class=\"Fz\(0\)\">).*?(</span>)", "", html)
-        df = (
-            pd.read_html(html_clean, header=None)[0]
-            .dropna(how="all", axis=1)
-            .fillna("-")
-            .replace("-", None)
-        )
+        df = pd.read_html(html_clean, header=None)[0].dropna(how="all", axis=1)
        return df

    @staticmethod
@@ -84,11 +81,12 @@ class YFGainersFetcher(Fetcher[YFGainersQueryParams, List[YFGainersData]]):
        **kwargs: Any,
    ) -> List[YFGainersData]:
        """Transform data."""
-        data["% Change"] = data["% Change"].str.replace("%", "").astype(float)
-        data["Volume"] = data["Volume"].str.replace("M", "").astype(float) * 1000000
-        data["Avg Vol (3 month)"] = (
-            data["Avg Vol (3 month)"].str.replace("M", "").astype(float) * 1000000
-        )
+
+        columns = ["Market Cap", "Avg Vol (3 month)", "Volume", "% Change"]
+
+        data = df_transform_numbers(data=data, columns=columns)
+        data = data.fillna("N/A").replace("N/A", None)
+
        return [
            YFGainersData.model_validate(d)
            for d in data.sort_values("% Change", ascending=False).to_dict(
--- a/openbb_platform/providers/yfinance/openbb_yfinance/models/growth_tech_equities.py
+++ b/openbb_platform/providers/yfinance/openbb_yfinance/models/growth_tech_equities.py
@@ -1,15 +1,18 @@
 """Yahoo Finance Asset Performance Growth Tech Equities Model."""

+# pylint: disable=unused-argument
+
 import re
 from typing import Any, Dict, List, Optional

 import pandas as pd
-import requests
 from openbb_core.provider.abstract.fetcher import Fetcher
 from openbb_core.provider.standard_models.equity_performance import (
    EquityPerformanceData,
    EquityPerformanceQueryParams,
 )
+from openbb_core.provider.utils.helpers import make_request
+from openbb_yfinance.utils.helpers import df_transform_numbers
 from pydantic import Field


@@ -35,10 +38,12 @@ class YFGrowthTechEquitiesData(EquityPerformanceData):
        "pe_ratio_ttm": "PE Ratio (TTM)",
    }

-    market_cap: str = Field(
+    market_cap: Optional[float] = Field(
+        default=None,
        description="Market Cap.",
    )
-    avg_volume_3_months: float = Field(
+    avg_volume_3_months: Optional[float] = Field(
+        default=None,
        description="Average volume over the last 3 months in millions.",
    )
    pe_ratio_ttm: Optional[float] = Field(
@@ -65,18 +70,12 @@ class YFGrowthTechEquitiesFetcher(
    ) -> pd.DataFrame:
        """Get data from YF."""
        headers = {"user_agent": "Mozilla/5.0"}
-        html = requests.get(
+        html = make_request(
            "https://finance.yahoo.com/screener/predefined/growth_technology_stocks",
            headers=headers,
-            timeout=10,
        ).text
        html_clean = re.sub(r"(<span class=\"Fz\(0\)\">).*?(</span>)", "", html)
-        df = (
-            pd.read_html(html_clean, header=None)[0]
-            .dropna(how="all", axis=1)
-            .fillna("-")
-            .replace("-", None)
-        )
+        df = pd.read_html(html_clean, header=None)[0].dropna(how="all", axis=1)
        return df

    @staticmethod
@@ -86,11 +85,12 @@ class YFGrowthTechEquitiesFetcher(
        **kwargs: Any,
    ) -> List[YFGrowthTechEquitiesData]:
        """Transform data."""
-        data["% Change"] = data["% Change"].str.replace("%", "")
-        data["Volume"] = data["Volume"].str.replace("M", "").astype(float) * 1000000
-        data["Avg Vol (3 month)"] = (
-            data["Avg Vol (3 month)"].str.replace("M", "").astype(float) * 1000000
-        )
-        data = data.to_dict(orient="records")
-        data = sorted(data, key=lambda d: d["Volume"], reverse=query.sort == "desc")
-        return [YFGrowthTechEquitiesData.model_validate(d) for d in data]
+
+        columns = ["Market Cap", "Avg Vol (3 month)", "Volume", "% Change"]
+
+        data = df_transform_numbers(data=data, columns=columns)
+        data = data.fillna("N/A").replace("N/A", None)
+
+        return [
+            YFGrowthTechEquitiesData.model_validate(d) for d in data.to_dict("records")
+        ]
--- a/openbb_platform/providers/yfinance/openbb_yfinance/models/losers.py
+++ b/openbb_platform/providers/yfinance/openbb_yfinance/models/losers.py
@@ -1,15 +1,18 @@
 """Yahoo Finance Asset Performance Losers Model."""

+# pylint: disable=unused-argument
+
 import re
 from typing import Any, Dict, List, Optional

 import pandas as pd
-import requests
 from openbb_core.provider.abstract.fetcher import Fetcher
 from openbb_core.provider.standard_models.equity_performance import (
    EquityPerformanceData,
    EquityPerformanceQueryParams,
 )
+from openbb_core.provider.utils.helpers import make_request
+from openbb_yfinance.utils.helpers import df_transform_numbers
 from pydantic import Field


@@ -35,10 +38,10 @@ class YFLosersData(EquityPerformanceData):
        "pe_ratio_ttm": "PE Ratio (TTM)",
    }

-    market_cap: str = Field(
+    market_cap: Optional[float] = Field(
        description="Market Cap.",
    )
-    avg_volume_3_months: float = Field(
+    avg_volume_3_months: Optional[float] = Field(
        description="Average volume over the last 3 months in millions.",
    )
    pe_ratio_ttm: Optional[float] = Field(
@@ -63,18 +66,12 @@ class YFLosersFetcher(Fetcher[YFLosersQueryParams, List[YFLosersData]]):
    ) -> pd.DataFrame:
        """Get data from YF."""
        headers = {"user_agent": "Mozilla/5.0"}
-        html = requests.get(
+        html = make_request(
            "https://finance.yahoo.com/screener/predefined/day_losers",
            headers=headers,
-            timeout=10,
        ).text
        html_clean = re.sub(r"(<span class=\"Fz\(0\)\">).*?(</span>)", "", html)
-        df = (
-            pd.read_html(html_clean, header=None)[0]
-            .dropna(how="all", axis=1)
-            .fillna("-")
-            .replace("-", None)
-        )
+        df = pd.read_html(html_clean, header=None)[0].dropna(how="all", axis=1)

        return df

@@ -85,11 +82,12 @@ class YFLosersFetcher(Fetcher[YFLosersQueryParams, List[YFLosersData]]):
        **kwargs: Any,
    ) -> List[YFLosersData]:
        """Transform data."""
-        data["% Change"] = data["% Change"].str.replace("%", "").astype(float)
-        data["Volume"] = data["Volume"].str.replace("M", "").astype(float) * 1000000
-        data["Avg Vol (3 month)"] = (
-            data["Avg Vol (3 month)"].str.replace("M", "").astype(float) * 1000000
-        )
+
+        columns = ["Market Cap", "Avg Vol (3 month)", "Volume", "% Change"]
+
+        data = df_transform_numbers(data=data, columns=columns)
+        data = data.fillna("N/A").replace("N/A", None)
+
        return [
            YFLosersData.model_validate(d)
            for d in data.sort_values("% Change", ascending=True).to_dict(
--- a/openbb_platform/providers/yfinance/openbb_yfinance/models/undervalued_growth_equities.py
+++ b/openbb_platform/providers/yfinance/openbb_yfinance/models/undervalued_growth_equities.py
@@ -1,15 +1,18 @@
 """Yahoo Finance Asset Undervalued Growth Tech Equities Model."""

+# pylint: disable=unused-argument
+
 import re
 from typing import Any, Dict, List, Optional

 import pandas as pd
-import requests
 from openbb_core.provider.abstract.fetcher import Fetcher
 from openbb_core.provider.standard_models.equity_performance import (
    EquityPerformanceData,
    EquityPerformanceQueryParams,
 )
+from openbb_core.provider.utils.helpers import make_request
+from openbb_yfinance.utils.helpers import df_transform_numbers
 from pydantic import Field


@@ -35,10 +38,12 @@ class YFUndervaluedGrowthEquitiesData(EquityPerformanceData):
        "pe_ratio_ttm": "PE Ratio (TTM)",
    }

-    market_cap: str = Field(
+    market_cap: Optional[float] = Field(
+        default=None,
        description="Market Cap.",
    )
-    avg_volume_3_months: float = Field(
+    avg_volume_3_months: Optional[float] = Field(
+        default=None,
        description="Average volume over the last 3 months in millions.",
    )
    pe_ratio_ttm: Optional[float] = Field(
@@ -69,18 +74,12 @@ class YFUndervaluedGrowthEquitiesFetcher(
    ) -> pd.DataFrame:
        """Get data from YF."""
        headers = {"user_agent": "Mozilla/5.0"}
-        html = requests.get(
+        html = make_request(
            "https://finance.yahoo.com/screener/predefined/undervalued_growth_stocks",
            headers=headers,
-            timeout=10,
        ).text
        html_clean = re.sub(r"(<span class=\"Fz\(0\)\">).*?(</span>)", "", html)
-        df = (
-            pd.read_html(html_clean, header=None)[0]
-            .dropna(how="all", axis=1)
-            .fillna("-")
-            .replace("-", None)
-        )
+        df = pd.read_html(html_clean, header=None)[0].dropna(how="all", axis=1)
        return df

    @staticmethod
@@ -90,11 +89,15 @@ class YFUndervaluedGrowthEquitiesFetcher(
        **kwargs: Any,
    ) -> List[YFUndervaluedGrowthEquitiesData]:
        """Transform data."""
-        data["% Change"] = data["% Change"].str.replace("%", "")
-        data["Volume"] = data["Volume"].str.replace("M", "").astype(float) * 1000000
-        data["Avg Vol (3 month)"] = (
-            data["Avg Vol (3 month)"].str.replace("M", "").astype(float) * 1000000
-        )
-        data = data.to_dict(orient="records")
-        data = sorted(data, key=lambda d: d["Volume"], reverse=query.sort == "desc")
-        return [YFUndervaluedGrowthEquitiesData.model_validate(d) for d in data]
+
+        columns = ["Market Cap", "Avg Vol (3 month)", "Volume", "% Change"]
+
+        data = df_transform_numbers(data, columns)
+        data = data.fillna("N/A").replace("N/A", None)
+
+        return [
+            YFUndervaluedGrowthEquitiesData.model_validate(d)
+            for d in data.sort_values(by="Market Cap", ascending=False).to_dict(
+                "records"
+            )
+        ]
--- a/openbb_platform/providers/yfinance/openbb_yfinance/models/undervalued_large_caps.py
+++ b/openbb_platform/providers/yfinance/openbb_yfinance/models/undervalued_large_caps.py
@@ -1,15 +1,18 @@
 """Yahoo Finance Asset Undervalued Large Caps Model."""

+# pylint: disable=unused-argument
+
 import re
 from typing import Any, Dict, List, Optional

 import pandas as pd
-import requests
 from openbb_core.provider.abstract.fetcher import Fetcher
 from openbb_core.provider.standard_models.equity_performance import (
    EquityPerformanceData,
    EquityPerformanceQueryParams,
 )
+from openbb_core.provider.utils.helpers import make_request
+from openbb_yfinance.utils.helpers import df_transform_numbers
 from pandas import DataFrame
 from pydantic import Field

@@ -36,10 +39,12 @@ class YFUndervaluedLargeCapsData(EquityPerformanceData):
        "pe_ratio_ttm": "PE Ratio (TTM)",
    }

-    market_cap: float = Field(
+    market_cap: Optional[float] = Field(
+        default=None,
        description="Market Cap.",
    )
-    avg_volume_3_months: float = Field(
+    avg_volume_3_months: Optional[float] = Field(
+        default=None,
        description="Average volume over the last 3 months in millions.",
    )
    pe_ratio_ttm: Optional[float] = Field(
@@ -66,18 +71,12 @@ class YFUndervaluedLargeCapsFetcher(
    ) -> DataFrame:
        """Get data from YF."""
        headers = {"user_agent": "Mozilla/5.0"}
-        html = requests.get(
+        html = make_request(
            "https://finance.yahoo.com/screener/predefined/undervalued_large_caps",
            headers=headers,
-            timeout=10,
        ).text
        html_clean = re.sub(r"(<span class=\"Fz\(0\)\">).*?(</span>)", "", html)
-        df = (
-            pd.read_html(html_clean, header=None)[0]
-            .dropna(how="all", axis=1)
-            .fillna("-")
-            .replace("-", None)
-        )
+        df = pd.read_html(html_clean, header=None)[0].dropna(how="all", axis=1)
        return df

    @staticmethod
@@ -87,11 +86,13 @@ class YFUndervaluedLargeCapsFetcher(
        **kwargs: Any,
    ) -> List[YFUndervaluedLargeCapsData]:
        """Transform data."""
-        data["% Change"] = data["% Change"].str.replace("%", "")
-        data["Volume"] = data["Volume"].str.replace("M", "").astype(float) * 1000000
-        data["Avg Vol (3 month)"] = (
-            data["Avg Vol (3 month)"].str.replace("M", "").astype(float) * 1000000
-        )
-        data["Market Cap"] = data["Market Cap"].str.replace("B", "").astype(float)
-        data = data.to_dict(orient="records")
-        return [YFUndervaluedLargeCapsData.model_validate(d) for d in data]
+
+        columns = ["Market Cap", "Avg Vol (3 month)", "Volume", "% Change"]
+
+        data = df_transform_numbers(data, columns)
+        data = data.fillna("N/A").replace("N/A", None)
+
+        return [
+            YFUndervaluedLargeCapsData.model_validate(d)
+            for d in data.to_dict("records")
+        ]
--- a/openbb_platform/providers/yfinance/openbb_yfinance/utils/helpers.py
+++ b/openbb_platform/providers/yfinance/openbb_yfinance/utils/helpers.py
@@ -190,3 +190,20 @@ def yf_download(
        data.columns = data.columns.str.lower().str.replace(" ", "_").to_list()

    return data
+
+
+def df_transform_numbers(data: pd.DataFrame, columns: list) -> pd.DataFrame:
+    """Replace abbreviations of numbers with actual numbers."""
+    multipliers = {"M": 1e6, "B": 1e9, "T": 1e12}
+
+    def replace_suffix(x, suffix, multiplier):
+        return float(str(x).replace(suffix, "")) * multiplier if suffix in str(x) else x
+
+    for col in columns:
+        if col == "% Change":
+            data[col] = data[col].astype(str).str.replace("%", "").astype(float)
+        else:
+            for suffix, multiplier in multipliers.items():
+                data[col] = data[col].apply(replace_suffix, args=(suffix, multiplier))
+
+    return data