[BugFix] Charting Integration Tests (OpenBB-finance#6569)

* fixing charting tests * lint
robert0825 · Jul 9, 2024 · 8ac72c6 · 8ac72c6
1 parent b116df2
commit 8ac72c6
Show file tree

Hide file tree

Showing 5 changed files with 23 additions and 18 deletions.
diff --git a/openbb_platform/core/openbb_core/app/model/charts/chart.py b/openbb_platform/core/openbb_core/app/model/charts/chart.py
@@ -1,26 +1,19 @@
 """OpenBB Core Chart model."""
 
-from enum import Enum
 from typing import Any, Dict, Optional
 
 from pydantic import BaseModel, ConfigDict, Field
 
 
-class ChartFormat(str, Enum):
-    """Chart format."""
-
-    plotly = "plotly"
-
-
 class Chart(BaseModel):
     """Model for Chart."""
 
     content: Optional[Dict[str, Any]] = Field(
         default=None,
         description="Raw textual representation of the chart.",
     )
-    format: Optional[ChartFormat] = Field(
-        default=ChartFormat.plotly,
+    format: Optional[str] = Field(
+        default=None,
         description="Complementary attribute to the `content` attribute. It specifies the format of the chart.",
     )
     fig: Optional[Any] = Field(

diff --git a/openbb_platform/core/tests/app/model/charts/test_chart.py b/openbb_platform/core/tests/app/model/charts/test_chart.py
@@ -1,7 +1,7 @@
 """Test the chart model."""
 
 import pytest
-from openbb_core.app.model.charts.chart import Chart, ChartFormat
+from openbb_core.app.model.charts.chart import Chart
 
 
 def test_charting_default_values():
@@ -11,14 +11,14 @@ def test_charting_default_values():
 
     # Assert
     assert chart.content is None
-    assert chart.format == ChartFormat.plotly
+    assert chart.format is None
 
 
 def test_charting_custom_values():
     """Test the charting custom values."""
     # Arrange
     content = {"data": [1, 2, 3]}
-    chart_format = ChartFormat.plotly
+    chart_format = "plotly"
 
     # Act
     chart = Chart(content=content, format=chart_format)
@@ -42,7 +42,7 @@ def test_charting_config_validation():
     """Test the charting config validation."""
     # Arrange
     content = {"data": [1, 2, 3]}
-    chart_format = ChartFormat.plotly
+    chart_format = "plotly"
 
     chart = Chart(content=content, format=chart_format)
 

diff --git a/openbb_platform/obbject_extensions/charting/integration/test_charting_api.py b/openbb_platform/obbject_extensions/charting/integration/test_charting_api.py
@@ -799,13 +799,17 @@ def test_charting_derivatives_futures_historical(params, headers):
         (
             {
                 "provider": "yfinance",
-                "symbol": "VX",
+                "symbol": "ES",
+                "date": None,
+                "chart": True,
             }
         ),
         (
             {
                 "provider": "cboe",
                 "symbol": "VX",
+                "date": "2024-06-25",
+                "chart": True,
             }
         ),
     ],

diff --git a/openbb_platform/obbject_extensions/charting/integration/test_charting_python.py b/openbb_platform/obbject_extensions/charting/integration/test_charting_python.py
@@ -654,14 +654,17 @@ def test_charting_derivatives_futures_historical(params, obb):
         (
             {
                 "provider": "yfinance",
-                "symbol": "VX",
+                "symbol": "ES",
+                "date": None,
+                "chart": True,
             }
         ),
         (
             {
                 "provider": "cboe",
                 "symbol": "VX",
                 "date": "2024-06-25",
+                "chart": True,
             }
         ),
     ],

diff --git a/openbb_platform/obbject_extensions/charting/openbb_charting/charting.py b/openbb_platform/obbject_extensions/charting/openbb_charting/charting.py
@@ -64,6 +64,7 @@ class Charting:
         entry_point.load()
         for entry_point in entry_points(group="openbb_charting_extension")
     ]
+    _format = "plotly"  # the charts computed by this extension will be in plotly format
 
     def __init__(self, obbject):
         """Initialize Charting extension."""
@@ -361,15 +362,17 @@ def show(self, render: bool = True, **kwargs):
             fig, content = charting_function(**kwargs)
             fig = self._set_chart_style(fig)
             content = fig.show(external=True, **kwargs).to_plotly_json()
-            self._obbject.chart = Chart(fig=fig, content=content)
+            self._obbject.chart = Chart(fig=fig, content=content, format=self._format)
             if render:
                 fig.show(**kwargs)
         except Exception:  # pylint: disable=W0718
             try:
                 fig = self.create_line_chart(data=self._obbject.results, render=False, **kwargs)  # type: ignore
                 fig = self._set_chart_style(fig)
                 content = fig.show(external=True, **kwargs).to_plotly_json()  # type: ignore
-                self._obbject.chart = Chart(fig=fig, content=content)
+                self._obbject.chart = Chart(
+                    fig=fig, content=content, format=self._format
+                )
                 if render:
                     return fig.show(**kwargs)  # type: ignore
             except Exception as e:
@@ -480,7 +483,9 @@ def to_chart(
                 fig = self.create_line_chart(data=data_as_df, render=False, **kwargs)
                 fig = self._set_chart_style(fig)
                 content = fig.show(external=True, **kwargs).to_plotly_json()  # type: ignore
-                self._obbject.chart = Chart(fig=fig, content=content)
+                self._obbject.chart = Chart(
+                    fig=fig, content=content, format=self._format
+                )
                 if render:
                     return fig.show(**kwargs)  # type: ignore
             except Exception as e:  # pylint: disable=W0718