test: testing sma

auto_trading/indicators/ema.py

@@ -25,4 +25,4 @@ class EMA(Indicator):
             DataFrame: Stock valuated float.
                 For each stock give -1 if realy bad and +1 if realy good.
         """
-        return data.ewm(self.nb_values).mean().close.loc[data.index[-1][0]]
+        return data.close.unstack().ewm(self.nb_values).mean().loc[data.index[-1][0]]

auto_trading/indicators/sma.py

@@ -25,4 +25,6 @@ class SMA(Indicator):
             DataFrame: Stock valuated float.
                 For each stock give -1 if realy bad and +1 if realy good.
         """
-        return data.rolling(self.nb_values).mean().close.loc[data.index[-1][0]]
+        return (
+            data.close.unstack().rolling(self.nb_values).mean().loc[data.index[-1][0]]
+        )

main.py

@@ -12,7 +12,7 @@ logging.basicConfig(level=logging.DEBUG)
 
 
 if __name__ == "__main__":
-    bt = Backtest("./data/NYSE.csv")
+    bt = Backtest("./data/NYSE_small.csv")
     ptf = InMemoryPortfolio(
         base_balance=100, change_rate_getter=lambda: bt.current_change
     )

tests/indicators/test_sma.py

@@ -0,0 +1,34 @@
+"""Realy basic tests."""
+import pytest
+from datetime import datetime
+from pandas import DataFrame, Series  # type: ignore
+
+from auto_trading.indicators.sma import SMA
+
+
+def date(j: int) -> datetime:
+    return datetime.strptime(f"2015-03-{j}", "%Y-%m-%d")
+
+
+@pytest.mark.parametrize(
+    "nb_values, data, expected",
+    [
+        (
+            10,
+            DataFrame(
+                {
+                    "close": {
+                        (date(i), st): 10 if i == 10 and st == "GOOG" else 0
+                        for i in range(1, 11)
+                        for st in ("GOOG", "GOOGL")
+                    }
+                }
+            ),
+            {"GOOG": 1.0, "GOOGL": 0.0},
+        ),
+    ],
+)
+def test_sma(nb_values: int, data: DataFrame, expected: Series) -> None:
+    """Test the SMA indicator."""
+    sma = SMA(nb_values)
+    assert sma(data).to_dict() == expected