Browse Source

test: test sma

Signed-off-by: QuentinN42 <quentin@lieumont.fr>
master
QuentinN42 7 months ago
parent
commit
7c5312f8ba
Signed by: number42 GPG Key ID: 2CD7D563712B3A50
  1. 6
      auto_trading/indicators/sma.py
  2. 4
      auto_trading/indicators/sma2.py
  3. 34
      tests/indicators/test_sma2.py

6
auto_trading/indicators/sma.py

@ -26,5 +26,9 @@ class SMA(Indicator):
For each stock give -1 if realy bad and +1 if realy good.
"""
return (
data.close.unstack().rolling(self.nb_values).mean().loc[data.index[-1][0]]
data.close.unstack()
.rolling(self.nb_values)
.mean()
.fillna(0)
.loc[data.index[-1][0]]
)

4
auto_trading/indicators/sma2.py

@ -1,5 +1,5 @@
"""SMA indicator, not ready yet."""
import pandas as pd
import pandas as pd # type: ignore
from ..interfaces import Indicator
@ -29,7 +29,7 @@ class SMA(Indicator):
data3 = data2["close"]
# apply sma
sma = data3.rolling(window=self.windowSize).mean()
sma = data3.rolling(window=self.windowSize).mean().fillna(0)
# computes slope at each point of the sma
smaDiff = sma.diff()

34
tests/indicators/test_sma2.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.sma2 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
Loading…
Cancel
Save