OpenBB/tests/openbb_terminal/forecast/test_forecast_controller.py

from typing import List
import argparse
import pandas as pd
import pytest
from prompt_toolkit.completion import NestedCompleter

try:
    from openbb_terminal.forecast import forecast_controller as fc
except ImportError:
    pytest.skip(allow_module_level=True)

# pylint: disable=E1121
base = "openbb_terminal.forecast.forecast_controller."
df = pd.DataFrame([[1, 2, 3], [2, 3, 4], [3, 4, 5]], columns=["first", "date", "close"])


class Thing:
    help = True

    def format_help(self):
        return "You need help"


def mock_func(_):
    return (Thing(), 2)


def test_check_greater_than_one():
    assert fc.check_greater_than_one(2) == 2


def test_check_greater_than_one_invalid():
    with pytest.raises(argparse.ArgumentTypeError):
        fc.check_greater_than_one(-2)


def test_forecast_controller(mocker):
    mocker.patch(base + "session", True)
    mocker.patch(base + "obbff.USE_PROMPT_TOOLKIT", True)
    cont = fc.ForecastController(data=df, ticker="TSLA")
    assert isinstance(cont.files, List)


def test_fc_update_runtime_choices(mocker):
    mocker.patch(base + "session", True)
    mocker.patch(base + "obbff.USE_PROMPT_TOOLKIT", True)
    cont = fc.ForecastController()
    cont.datasets = {"stonks": df}
    cont.update_runtime_choices()
    assert isinstance(cont.completer, NestedCompleter)


def test_fc_refresh_datasets_on_menu():
    cont = fc.ForecastController()
    cont.datasets = {"stonks": df}
    cont.files = ["file.csv", "fiile.csv"]
    cont.refresh_datasets_on_menu()
    assert "stonks.first" in cont.list_dataset_cols


def test_fc_print_help(capsys):
    cont = fc.ForecastController()
    cont.print_help()
    captured = capsys.readouterr()
    assert "forecast" in captured.out


def test_fc_custom_reset():
    cont = fc.ForecastController()
    val = cont.custom_reset()
    assert not val


def test_fc_custom_reset_with_files():
    cont = fc.ForecastController()
    cont.files_full = [["file1", "file1"]]
    val = cont.custom_reset()
    assert val == ["forecast", "'load file1 -a file1'"]


def test_fc_parse_known_args_and_warn(mocker):
    mock = mocker.Mock()
    mock2 = mocker.Mock()
    mock.parse_known_args = mock_func
    mock.format_help = Thing().format_help
    mock.add_argument = mock2
    cont = fc.ForecastController()
    cont.parse_known_args_and_warn(
        parser=mock,
        other_args=[],
        export_allowed=0,
        raw=True,
        limit=1,
        target_dataset=True,
        target_column=True,
        period=5,
        n_days=True,
        seasonal="weekly",
        periods=True,
        window=True,
        train_split=True,
        input_chunk_length=True,
        output_chunk_length=True,
        force_reset=True,
        save_checkpoints=True,
        model_save_name="the_models",
        n_epochs=True,
        model_type=True,
        dropout=0.05,
        batch_size=3,
        learning_rate=True,
        past_covariates=True,
        lags=True,
        hidden_size=4,
        n_jumps=True,
        end=True,
        residuals=True,
        forecast_only=True,
    )
    assert mock2.call_count == 28


def test_fc_call_load(mocker):
    mocker.patch(base + "forecast_model.load", return_value=df)
    cont = fc.ForecastController()
    cont.call_load(["data.csv"])


def test_fc_call_load_alias():
    cont = fc.ForecastController()
    cont.files = ["the_data"]
    cont.call_load(["data.csv", "-a", "the_data"])


def test_fc_call_load_alias_no_dot():
    cont = fc.ForecastController()
    cont.call_load(["datacsv", "-a", "the_data"])


@pytest.mark.parametrize(
    "name,loc_df", [("base", df), ("base", pd.DataFrame()), ("", df)]
)
def test_fc_call_show(mocker, name, loc_df):
    mock = mocker.MagicMock()
    mock.name = name
    mock.limit = 4
    mocker.patch(
        base + "ForecastController.parse_known_args_and_warn", return_value=mock
    )
    cont = fc.ForecastController()
    cont.datasets = {"base": loc_df}
    cont.call_show(["data"])


@pytest.mark.parametrize("dataset", ["base", None])
def test_call_desc(mocker, dataset):
    mock = mocker.MagicMock()
    mock.target_dataset = dataset
    mocker.patch(
        base + "ForecastController.parse_known_args_and_warn", return_value=mock
    )
    cont = fc.ForecastController()
    cont.datasets = {"base": df}
    cont.call_desc(["data"])


def test_call_plot(mocker):
    mock = mocker.MagicMock()
    mock.values = ["data.first"]
    mocker.patch(
        base + "ForecastController.parse_known_args_and_warn", return_value=mock
    )
    cont = fc.ForecastController()
    cont.choices = {"plot": ["data.first", "data.second", "data.third"]}
    cont.datasets = {"data": df}
    cont.call_plot(["data.first"])


def test_call_plot_no_parser(mocker):
    mocker.patch(
        base + "ForecastController.parse_known_args_and_warn", return_value=None
    )
    cont = fc.ForecastController()
    cont.choices = {"plot": ["data.first", "data.second", "data.third"]}
    cont.datasets = {"data": df}
    cont.call_plot(["data.first"])


def test_call_plot_no_values(mocker):
    mock = mocker.MagicMock()
    mock.values = None
    mocker.patch(
        base + "ForecastController.parse_known_args_and_warn", return_value=mock
    )
    cont = fc.ForecastController()
    cont.choices = {"plot": ["data.first", "data.second", "data.third"]}
    cont.datasets = {"data": df}
    cont.call_plot(["data.first"])


def test_call_season(tsla_csv):
    cont = fc.ForecastController()
    cont.datasets = {"data": tsla_csv}
    cont.call_season(["data.close", "--max_lag", "1"])


def test_call_season_error():
    cont = fc.ForecastController()
    cont.datasets = {"data": df}
    cont.call_season(["datafirst", "--max_lag", "1"])


def test_call_season_no_parser(mocker):
    mocker.patch(
        base + "ForecastController.parse_known_args_and_warn", return_value=None
    )
    cont = fc.ForecastController()
    cont.datasets = {"data": df}
    cont.call_season(["data.first", "--max_lag", "1"])


def test_call_season_no_values(mocker):
    mock = mocker.MagicMock()
    mock.values = None
    mocker.patch(
        base + "ForecastController.parse_known_args_and_warn", return_value=mock
    )
    cont = fc.ForecastController()
    cont.datasets = {"data": df}
    cont.call_season(["data.first", "--max_lag", "1"])


def test_call_corr():
    cont = fc.ForecastController()
    cont.datasets = {"data": df}
    cont.call_corr(["data"])


def test_call_corr_target(mocker):
    mock = mocker.MagicMock()
    mock.target_dataset = None
    mocker.patch(
        base + "ForecastController.parse_known_args_and_warn", return_value=mock
    )
    cont = fc.ForecastController()
    cont.datasets = {"data": df}
    cont.call_corr(["data"])


def test_call_comb_not_in(mocker):
    mock = mocker.MagicMock()
    mock.values = ["data.first"]
    mocker.patch(
        base + "ForecastController.parse_known_args_and_warn", return_value=mock
    )
    cont = fc.ForecastController()
    cont.datasets = {"data": df}
    cont.choices = {
        "combine": {"data.first": 1, "data.second": 2, "data.third": 3},
        "delete": {"data.first": 1, "data.second": 2, "data.third": 3},
    }
    cont.call_combine(["data"])


def test_call_comb(mocker):
    mock = mocker.MagicMock()
    mock.dataset = "data"
    mock.columns = ["data.first", "data.second"]
    mocker.patch(
        base + "ForecastController.parse_known_args_and_warn", return_value=mock
    )
    cont = fc.ForecastController()
    cont.datasets = {"data": df}
    cont.choices = {
        "combine": {"data.first": 1, "data.second": 2, "data.third": 3},
        "delete": {"data.first": 1, "data.second": 2, "data.third": 3},
    }
    cont.call_combine(["data"])


def test_call_clean():
    cont = fc.ForecastController()
    cont.datasets = {"data": df}
    cont.call_clean(["data"])


def test_call_clean_bad(mocker):
    mock = mocker.MagicMock()
    mock.target_dataset = None
    mocker.patch(
        base + "ForecastController.parse_known_args_and_warn", return_value=mock
    )
    cont = fc.ForecastController()
    cont.datasets = {"data": df}
    cont.call_clean(["data"])


@pytest.mark.parametrize(
    "feature", ["ema", "sto", "rsi", "roc", "mom", "delta", "atr", "signal", "delete"]
)
def test_call_feat_eng_invalid(feature):
    cont = fc.ForecastController()
    cont.datasets = {"data": df}
    cont.files = ["file.csv", "fiile.csv"]
    cont.choices = {
        "combine": {"data.first": 1, "data.second": 2, "data.third": 3},
        "delete": {"data.first": 1, "data.second": 2, "data.third": 3},
    }
    the_list = ["data"]
    if feature == "rsi":
        the_list.append("--period")
        the_list.append("2")
    getattr(cont, f"call_{feature}")(the_list)


@pytest.mark.parametrize(
    "feature",
    ["ema", "sto", "rsi", "roc", "mom", "delta", "atr", "signal", "delete", "export"],
)
def test_call_feat_eng_invalid_parser(feature, mocker):
    mocker.patch(base + "helpers.check_parser_input", return_value=None)
    mock = mocker.MagicMock()
    mock.target_dataset = None
    mocker.patch(
        base + "ForecastController.parse_known_args_and_warn", return_value=mock
    )
    cont = fc.ForecastController()
    cont.datasets = {"data": df}
    cont.files = ["file.csv", "fiile.csv"]
    cont.choices = {
        "combine": {"data.first": 1, "data.second": 2, "data.third": 3},
        "delete": {"data.first": 1, "data.second": 2, "data.third": 3},
    }
    the_list = ["data"]
    if feature == "rsi":
        the_list.append("--period")
        the_list.append("2")
    getattr(cont, f"call_{feature}")(the_list)


def test_call_ema(mocker):
    mock = mocker.MagicMock()
    mock.target_dataset = "data"
    mock.columns = ["data.first", "data.second"]
    mocker.patch(
        base + "ForecastController.parse_known_args_and_warn", return_value=mock
    )
    cont = fc.ForecastController()
    cont.datasets = {"data": df}
    cont.call_ema(["data"])


@pytest.mark.parametrize("datasets", [[], ["data"], ["bad"]])
def test_call_delete(mocker, datasets):
    cont = fc.ForecastController()
    mock = mocker.MagicMock()
    if datasets == ["bad"]:
        mock.delete = ["data.lose"]
    else:
        mock.delete = ["data.close"]
    mocker.patch(
        base + "ForecastController.parse_known_args_and_warn", return_value=mock
    )
    if datasets == ["data"]:
        cont.datasets = {"data": df}
    cont.choices = {
        "combine": ["data.first", "data.second", "data.third"],
        "delete": ["data.first", "data.second", "data.third"],
    }
    cont.call_delete([])


def test_call_export(mocker):
    mock = mocker.MagicMock()
    mock.target_dataset = "data"
    mock.columns = ["data.first", "data.second"]
    mocker.patch(
        base + "ForecastController.parse_known_args_and_warn", return_value=mock
    )
    cont = fc.ForecastController()
    cont.datasets = {"data": df}
    cont.call_export(["data"])


@pytest.mark.parametrize(
    "opt, func",
    [
        ("expo", "expo_view.display_expo_forecast"),
        ("theta", "theta_view.display_theta_forecast"),
        ("rnn", "rnn_view.display_rnn_forecast"),
        ("nbeats", "nbeats_view.display_nbeats_forecast"),
        ("tcn", "tcn_view.display_tcn_forecast"),
        ("regr", "regr_view.display_regression"),
        ("linregr", "linregr_view.display_linear_regression"),
        ("brnn", "brnn_view.display_brnn_forecast"),
        ("trans", "trans_view.display_trans_forecast"),
        ("tft", "tft_view.display_tft_forecast"),
    ],
)
def test_models(mocker, opt, func):
    mocker.patch(base + "helpers.check_parser_input", return_value=True)
    mocker.patch(base + func)
    cont = fc.ForecastController()
    cont.datasets = {"data": df}
    getattr(cont, f"call_{opt}")(["data"])


@pytest.mark.parametrize(
    "opt",
    [
        "expo",
        "theta",
        "rnn",
        "nbeats",
        "tcn",
        "regr",
        "linregr",
        "brnn",
        "trans",
        "tft",
    ],
)
def test_models_bad(opt):
    cont = fc.ForecastController()
    cont.datasets = {"data": df}
    getattr(cont, f"call_{opt}")(["data"])