From 75ca285a50a4fb79538be4e12b37def7efd46140 Mon Sep 17 00:00:00 2001
From: "d.a.bunin" <d.a.bunin@tinkoff.ru>
Date: Fri, 2 Sep 2022 18:51:15 +0300
Subject: [PATCH 1/9] Make SeasonalMovingAverageModel to work with context,
 make DeadlineMovingAverageModel to work with context

---
 etna/models/base.py                     | 13 ++++-
 etna/models/deadline_ma.py              | 77 +++++++++++++++----------
 etna/models/seasonal_ma.py              | 48 ++++++++-------
 tests/test_models/test_simple_models.py | 44 +++++++++-----
 4 files changed, 115 insertions(+), 67 deletions(-)

diff --git a/etna/models/base.py b/etna/models/base.py
index f776c20d9..2e0cc8c50 100644
--- a/etna/models/base.py
+++ b/etna/models/base.py
@@ -462,7 +462,12 @@ def _forecast_segment(model: Any, segment: str, ts: TSDataset, *args, **kwargs)
         if isinstance(segment_predict, np.ndarray):
             segment_predict = pd.DataFrame({"target": segment_predict})
         segment_predict["segment"] = segment
-        segment_predict["timestamp"] = dates
+
+        prediction_size = kwargs.get("prediction_size")
+        if prediction_size is not None:
+            segment_predict["timestamp"] = dates[-prediction_size:].reset_index(drop=True)
+        else:
+            segment_predict["timestamp"] = dates
         return segment_predict
 
     @log_decorator
@@ -489,11 +494,16 @@ def _forecast(self, ts: TSDataset, **kwargs) -> TSDataset:
         result_df = result_df.set_index(["timestamp", "segment"])
         df = ts.to_pandas(flatten=True)
         df = df.set_index(["timestamp", "segment"])
+        # TODO: remember that it can be a trouble for in-sample forecasting
         df = df.combine_first(result_df).reset_index()
 
         df = TSDataset.to_dataset(df)
         ts.df = df
         ts.inverse_transform()
+
+        prediction_size = kwargs.get("prediction_size")
+        if prediction_size is not None:
+            ts.df = ts.df.iloc[-prediction_size:]
         return ts
 
 
@@ -547,6 +557,7 @@ def _forecast(self, ts: TSDataset, **kwargs) -> TSDataset:
         """
         horizon = len(ts.df)
         x = ts.to_pandas(flatten=True).drop(["segment"], axis=1)
+        # TODO: we haven't tested working with prediction_size here
         y = self._base_model.predict(x, **kwargs).reshape(-1, horizon).T
         ts.loc[:, pd.IndexSlice[:, "target"]] = y
         ts.inverse_transform()
diff --git a/etna/models/deadline_ma.py b/etna/models/deadline_ma.py
index c2610bbd7..bde8e99d2 100644
--- a/etna/models/deadline_ma.py
+++ b/etna/models/deadline_ma.py
@@ -6,8 +6,8 @@
 import numpy as np
 import pandas as pd
 
-from etna.models.base import NonPredictionIntervalContextIgnorantAbstractModel
-from etna.models.base import NonPredictionIntervalContextIgnorantModelMixin
+from etna.models.base import NonPredictionIntervalContextRequiredAbstractModel
+from etna.models.base import NonPredictionIntervalContextRequiredModelMixin
 from etna.models.base import PerSegmentModelMixin
 
 
@@ -31,7 +31,7 @@ def __init__(self, window: int = 3, seasonality: str = "month"):
         """
         Initialize deadline moving average model.
 
-        Length of remembered tail of series is equal to the number of ``window`` months or years, depending on the ``seasonality``.
+        Length of the context is equal to the number of ``window`` months or years, depending on the ``seasonality``.
 
         Parameters
         ----------
@@ -78,28 +78,30 @@ def fit(self, df: pd.DataFrame, regressors: List[str]) -> "_DeadlineMovingAverag
                 message=f"{type(self).__name__} does not work with any exogenous series or features. "
                 f"It uses only target series for predict/\n "
             )
-        targets = df["target"]
-        timestamps = df["timestamp"]
+
+        self._freq = freq
+
+        return self
+
+    def _get_context_beginning(self, df: pd.DataFrame, prediction_size: int):
+        df_history = df.iloc[:-prediction_size]
+        history_timestamps = df_history["timestamp"]
+        future_timestamps = df["timestamp"].iloc[-prediction_size:]
 
         if self.seasonality == SeasonalityMode.month:
-            first_index = timestamps.iloc[-1] - pd.DateOffset(months=self.window)
+            first_index = future_timestamps.iloc[0] - pd.DateOffset(months=self.window)
 
         elif self.seasonality == SeasonalityMode.year:
-            first_index = timestamps.iloc[-1] - pd.DateOffset(years=self.window)
+            first_index = future_timestamps.iloc[0] - pd.DateOffset(years=self.window)
 
-        if first_index < timestamps.iloc[0]:
+        if first_index < history_timestamps.iloc[0]:
             raise ValueError(
-                "Given series is too short for chosen shift value. Try lower shift value, or give" "longer series."
+                "Given context isn't big enough, try to decrease context_size, prediction_size of increase length of given dataframe!"
             )
 
-        self.series = targets.loc[timestamps >= first_index]
-        self.timestamps = timestamps.loc[timestamps >= first_index]
-        self.shift = len(self.series)
-        self._freq = freq
-
-        return self
+        return first_index
 
-    def predict(self, df: pd.DataFrame) -> np.ndarray:
+    def predict(self, df: pd.DataFrame, prediction_size: int) -> np.ndarray:
         """
         Compute predictions from a DeadlineMovingAverageModel.
 
@@ -107,36 +109,49 @@ def predict(self, df: pd.DataFrame) -> np.ndarray:
         ----------
         df: pd.DataFrame
             Used only for getting the horizon of forecast and timestamps.
+        prediction_size:
+            Number of last timestamps to leave after making prediction.
+            Previous timestamps will be used as a context for models that require it.
 
         Returns
         -------
         :
             Array with predictions.
+
+        Raises
+        ------
+        ValueError:
+            if context isn't big enought
         """
-        timestamps = df["timestamp"]
-        index = pd.date_range(start=self.timestamps.iloc[0], end=timestamps.iloc[-1])
-        res = np.append(self.series.values, np.zeros(len(df)))
+        context_beginning = self._get_context_beginning(df=df, prediction_size=prediction_size)
+
+        df_history = df.iloc[:-prediction_size]
+        history_targets = df_history["target"]
+        history_timestamps = df_history["timestamp"]
+        future_timestamps = df["timestamp"].iloc[-prediction_size:]
+        history_targets = history_targets.loc[history_timestamps >= context_beginning]
+        history_timestamps = history_timestamps.loc[history_timestamps >= context_beginning]
+
+        index = pd.date_range(start=context_beginning, end=future_timestamps.iloc[-1])
+        res = np.append(history_targets.values, np.zeros(prediction_size))
         res = pd.DataFrame(res)
         res.index = index
-        for i in range(len(self.series), len(res)):
+        for i in range(len(history_targets), len(res)):
             for w in range(1, self.window + 1):
                 if self.seasonality == SeasonalityMode.month:
                     prev_date = res.index[i] - pd.DateOffset(months=w)
-
                 elif self.seasonality == SeasonalityMode.year:
                     prev_date = res.index[i] - pd.DateOffset(years=w)
-                if prev_date <= self.timestamps.iloc[-1]:
-                    res.loc[index[i]] += self.series.loc[self.timestamps == prev_date].values
+
+                if prev_date <= history_timestamps.iloc[-1]:
+                    res.loc[index[i]] += history_targets.loc[history_timestamps == prev_date].values
                 else:
                     res.loc[index[i]] += res.loc[prev_date].values
 
             res.loc[index[i]] = res.loc[index[i]] / self.window
 
-        res = res.values.reshape(
-            len(res),
-        )
-
-        return res[-len(df) :]
+        res = res.values.ravel()[-prediction_size:]
+        return res
 
     @property
     def context_size(self) -> int:
@@ -159,8 +174,8 @@ def context_size(self) -> int:
 
 class DeadlineMovingAverageModel(
     PerSegmentModelMixin,
-    NonPredictionIntervalContextIgnorantModelMixin,
-    NonPredictionIntervalContextIgnorantAbstractModel,
+    NonPredictionIntervalContextRequiredModelMixin,
+    NonPredictionIntervalContextRequiredAbstractModel,
 ):
     """Moving average model that uses exact previous dates to predict."""
 
@@ -168,6 +183,8 @@ def __init__(self, window: int = 3, seasonality: str = "month"):
         """
         Initialize deadline moving average model.
 
+        Length of the context is equal to the number of ``window`` months or years, depending on the ``seasonality``.
+
         Parameters
         ----------
         window: int
diff --git a/etna/models/seasonal_ma.py b/etna/models/seasonal_ma.py
index 6a989062f..87b08684d 100644
--- a/etna/models/seasonal_ma.py
+++ b/etna/models/seasonal_ma.py
@@ -5,8 +5,8 @@
 import numpy as np
 import pandas as pd
 
-from etna.models.base import NonPredictionIntervalContextIgnorantAbstractModel
-from etna.models.base import NonPredictionIntervalContextIgnorantModelMixin
+from etna.models.base import NonPredictionIntervalContextRequiredAbstractModel
+from etna.models.base import NonPredictionIntervalContextRequiredModelMixin
 from etna.models.base import PerSegmentModelMixin
 
 
@@ -24,7 +24,7 @@ def __init__(self, window: int = 5, seasonality: int = 7):
         """
         Initialize seasonal moving average model.
 
-        Length of remembered tail of series is ``window * seasonality``.
+        Length of the context is ``window * seasonality``.
 
         Parameters
         ----------
@@ -33,7 +33,6 @@ def __init__(self, window: int = 5, seasonality: int = 7):
         seasonality: int
             Lag between values taken for forecast.
         """
-        self.series = None
         self.name = "target"
         self.window = window
         self.seasonality = seasonality
@@ -45,7 +44,7 @@ def fit(self, df: pd.DataFrame, regressors: List[str]) -> "_SeasonalMovingAverag
 
         Parameters
         ----------
-        df: pd.DataFrame
+        df:
             Data to fit on
         regressors:
             List of the columns with regressors(ignored in this model)
@@ -60,44 +59,49 @@ def fit(self, df: pd.DataFrame, regressors: List[str]) -> "_SeasonalMovingAverag
                 message=f"{type(self).__name__} does not work with any exogenous series or features. "
                 f"It uses only target series for predict/\n "
             )
-        targets = df["target"]
-        if len(targets) < self.shift:
-            raise ValueError(
-                "Given series is too short for chosen shift value. Try lower shift value, or give" "longer series."
-            )
-        self.series = targets[-self.shift :].values
 
-        # ???
-        if targets.name is not None:
-            self.name = targets.name
         return self
 
-    def predict(self, df: pd.DataFrame) -> np.ndarray:
+    def predict(self, df: pd.DataFrame, prediction_size: int) -> np.ndarray:
         """
         Compute predictions from a SeasonalMovingAverage model.
 
         Parameters
         ----------
-        df: pd.DataFrame
+        df:
             Used only for getting the horizon of forecast
+        prediction_size:
+            Number of last timestamps to leave after making prediction.
+            Previous timestamps will be used as a context for models that require it.
 
         Returns
         -------
         :
             Array with predictions.
+
+        Raises
+        ------
+        ValueError:
+            if context isn't big enought
         """
-        horizon = len(df)
-        res = np.append(self.series, np.zeros(horizon))
+        expected_length = prediction_size + self.shift
+        if len(df) < expected_length:
+            raise ValueError(
+                "Given context isn't big enough, try to decrease context_size, prediction_size of increase length of given dataframe!"
+            )
+
+        history = df["target"][-expected_length:-prediction_size]
+        res = np.append(history, np.zeros(prediction_size))
         for i in range(self.shift, len(res)):
             res[i] = res[i - self.shift : i : self.seasonality].mean()
-        y_pred = res[-horizon:]
+        y_pred = res[-prediction_size:]
         return y_pred
 
 
 class SeasonalMovingAverageModel(
     PerSegmentModelMixin,
-    NonPredictionIntervalContextIgnorantModelMixin,
-    NonPredictionIntervalContextIgnorantAbstractModel,
+    NonPredictionIntervalContextRequiredModelMixin,
+    NonPredictionIntervalContextRequiredAbstractModel,
 ):
     """
     Seasonal moving average.
@@ -112,7 +116,7 @@ def __init__(self, window: int = 5, seasonality: int = 7):
         """
         Initialize seasonal moving average model.
 
-        Length of remembered tail of series is ``window * seasonality``.
+        Length of the context is ``window * seasonality``.
 
         Parameters
         ----------
diff --git a/tests/test_models/test_simple_models.py b/tests/test_models/test_simple_models.py
index 0e7fcf031..5fb9f77a7 100644
--- a/tests/test_models/test_simple_models.py
+++ b/tests/test_models/test_simple_models.py
@@ -39,29 +39,45 @@ def df():
 def test_simple_model_forecaster_run(simple_df, model):
     sma_model = model()
     sma_model.fit(simple_df)
-    future_ts = simple_df.make_future(future_steps=7)
-    res = sma_model.forecast(future_ts)
+    future_ts = simple_df.make_future(future_steps=7, tail_steps=sma_model.context_size)
+    res = sma_model.forecast(future_ts, prediction_size=7)
     res = res.to_pandas(flatten=True)
     assert not res.isnull().values.any()
     assert len(res) == 14
 
 
+def test_simple_model_forecaster_fail(simple_df):
+    sma_model = SeasonalMovingAverageModel(window=1000, seasonality=7)
+    sma_model.fit(simple_df)
+    future_ts = simple_df.make_future(future_steps=7, tail_steps=sma_model.context_size)
+    with pytest.raises(ValueError, match="Given context isn't big enough"):
+        _ = sma_model.forecast(future_ts, prediction_size=7)
+
+
 @pytest.mark.parametrize("model", [DeadlineMovingAverageModel])
 def test_deadline_model_forecaster_run(simple_df, model):
     model = model(window=1)
     model.fit(simple_df)
-    future_ts = simple_df.make_future(future_steps=7)
-    res = model.forecast(future_ts)
+    future_ts = simple_df.make_future(future_steps=7, tail_steps=model.context_size)
+    res = model.forecast(future_ts, prediction_size=7)
     res = res.to_pandas(flatten=True)
     assert not res.isnull().values.any()
     assert len(res) == 14
 
 
+def test_sdeadline_model_forecaster_fail(simple_df):
+    model = DeadlineMovingAverageModel(window=1000)
+    model.fit(simple_df)
+    future_ts = simple_df.make_future(future_steps=7, tail_steps=model.context_size)
+    with pytest.raises(ValueError, match="Given context isn't big enough"):
+        _ = model.forecast(future_ts, prediction_size=7)
+
+
 def test_seasonal_moving_average_forecaster_correct(simple_df):
     model = SeasonalMovingAverageModel(window=3, seasonality=7)
     model.fit(simple_df)
-    future_ts = simple_df.make_future(future_steps=7)
-    res = model.forecast(future_ts)
+    future_ts = simple_df.make_future(future_steps=7, tail_steps=model.context_size)
+    res = model.forecast(future_ts, prediction_size=7)
     res = res.to_pandas(flatten=True)[["target", "segment", "timestamp"]]
 
     df1 = pd.DataFrame()
@@ -83,8 +99,8 @@ def test_seasonal_moving_average_forecaster_correct(simple_df):
 def test_naive_forecaster_correct(simple_df):
     model = NaiveModel(lag=3)
     model.fit(simple_df)
-    future_ts = simple_df.make_future(future_steps=7)
-    res = model.forecast(future_ts)
+    future_ts = simple_df.make_future(future_steps=7, tail_steps=model.context_size)
+    res = model.forecast(future_ts, prediction_size=7)
     res = res.to_pandas(flatten=True)[["target", "segment", "timestamp"]]
 
     df1 = pd.DataFrame()
@@ -107,8 +123,8 @@ def test_naive_forecaster_correct(simple_df):
 def test_moving_average_forecaster_correct(simple_df):
     model = MovingAverageModel(window=5)
     model.fit(simple_df)
-    future_ts = simple_df.make_future(future_steps=7)
-    res = model.forecast(future_ts)
+    future_ts = simple_df.make_future(future_steps=7, tail_steps=model.context_size)
+    res = model.forecast(future_ts, prediction_size=7)
     res = res.to_pandas(flatten=True)[["target", "segment", "timestamp"]]
 
     df1 = pd.DataFrame()
@@ -137,8 +153,8 @@ def test_moving_average_forecaster_correct(simple_df):
 def test_deadline_moving_average_forecaster_correct(df):
     model = DeadlineMovingAverageModel(window=3, seasonality="month")
     model.fit(df)
-    future_ts = df.make_future(future_steps=20)
-    res = model.forecast(future_ts)
+    future_ts = df.make_future(future_steps=20, tail_steps=model.context_size)
+    res = model.forecast(future_ts, prediction_size=20)
     res = res.to_pandas(flatten=True)[["target", "segment", "timestamp"]]
 
     df1 = pd.DataFrame()
@@ -323,8 +339,8 @@ def two_month_ts():
 def test_deadline_model_correct_with_big_horizons(two_month_ts):
     model = DeadlineMovingAverageModel(window=2, seasonality="month")
     model.fit(two_month_ts)
-    future_ts = two_month_ts.make_future(future_steps=90)
-    res = model.forecast(future_ts)
+    future_ts = two_month_ts.make_future(future_steps=90, tail_steps=model.context_size)
+    res = model.forecast(future_ts, prediction_size=90)
     expected = np.array(
         [
             [16.5],

From da6480647900437554bc2d70adfd12c7a5fef57c Mon Sep 17 00:00:00 2001
From: "d.a.bunin" <d.a.bunin@tinkoff.ru>
Date: Mon, 5 Sep 2022 12:35:38 +0300
Subject: [PATCH 2/9] Fix test_inference tests

---
 etna/models/deadline_ma.py          |   2 +-
 tests/test_models/test_inference.py | 148 ++++++++++++++++++++--------
 2 files changed, 106 insertions(+), 44 deletions(-)

diff --git a/etna/models/deadline_ma.py b/etna/models/deadline_ma.py
index bde8e99d2..b193c51ab 100644
--- a/etna/models/deadline_ma.py
+++ b/etna/models/deadline_ma.py
@@ -94,7 +94,7 @@ def _get_context_beginning(self, df: pd.DataFrame, prediction_size: int):
         elif self.seasonality == SeasonalityMode.year:
             first_index = future_timestamps.iloc[0] - pd.DateOffset(years=self.window)
 
-        if first_index < history_timestamps.iloc[0]:
+        if len(history_timestamps) == 0 or first_index < history_timestamps.iloc[0]:
             raise ValueError(
                 "Given context isn't big enough, try to decrease context_size, prediction_size of increase length of given dataframe!"
             )
diff --git a/tests/test_models/test_inference.py b/tests/test_models/test_inference.py
index 7f43a484d..17498b7db 100644
--- a/tests/test_models/test_inference.py
+++ b/tests/test_models/test_inference.py
@@ -1,14 +1,19 @@
+from copy import deepcopy
+
 import numpy as np
 import pandas as pd
 import pytest
 from pandas.util.testing import assert_frame_equal
 from pytorch_forecasting.data import GroupNormalizer
+from typing_extensions import get_args
 
 from etna.datasets import TSDataset
 from etna.models import AutoARIMAModel
 from etna.models import BATSModel
 from etna.models import CatBoostModelMultiSegment
 from etna.models import CatBoostModelPerSegment
+from etna.models import ContextRequiredModelType
+from etna.models import DeadlineMovingAverageModel
 from etna.models import ElasticMultiSegmentModel
 from etna.models import ElasticPerSegmentModel
 from etna.models import HoltModel
@@ -39,14 +44,19 @@ def _test_forecast_in_sample_full(ts, model, transforms):
     forecast_ts = TSDataset(df, freq="D")
     forecast_ts.transform(ts.transforms)
     forecast_ts.df.loc[:, pd.IndexSlice[:, "target"]] = np.NaN
-    model.forecast(forecast_ts)
+
+    if isinstance(model, get_args(ContextRequiredModelType)):
+        prediction_size = len(forecast_ts.index)
+        model.forecast(forecast_ts, prediction_size=prediction_size)
+    else:
+        model.forecast(forecast_ts)
 
     # checking
     forecast_df = forecast_ts.to_pandas(flatten=True)
     assert not np.any(forecast_df["target"].isna())
 
 
-def _test_forecast_in_sample_suffix(ts, model, transforms):
+def _test_forecast_in_sample_suffix(ts, model, transforms, num_skip_points):
     df = ts.to_pandas()
 
     # fitting
@@ -56,9 +66,15 @@ def _test_forecast_in_sample_suffix(ts, model, transforms):
     # forecasting
     forecast_ts = TSDataset(df, freq="D")
     forecast_ts.transform(ts.transforms)
-    forecast_ts.df.loc[:, pd.IndexSlice[:, "target"]] = np.NaN
-    forecast_ts.df = forecast_ts.df.iloc[6:]
-    model.forecast(forecast_ts)
+
+    if isinstance(model, get_args(ContextRequiredModelType)):
+        prediction_size = len(forecast_ts.index) - num_skip_points
+        forecast_ts.df.loc[forecast_ts.index[num_skip_points] :, pd.IndexSlice[:, "target"]] = np.NaN
+        model.forecast(forecast_ts, prediction_size=prediction_size)
+    else:
+        forecast_ts.df = forecast_ts.df.iloc[num_skip_points:]
+        forecast_ts.df.loc[:, pd.IndexSlice[:, "target"]] = np.NaN
+        model.forecast(forecast_ts)
 
     # checking
     forecast_df = forecast_ts.to_pandas(flatten=True)
@@ -69,21 +85,30 @@ def _test_forecast_out_sample_prefix(ts, model, transforms):
     # fitting
     ts.fit_transform(transforms)
     model.fit(ts)
-    # forecasting full
-    forecast_full_ts = ts.make_future(5)
 
+    # forecasting full
     import torch  # TODO: remove after fix at issue-802
 
     torch.manual_seed(11)
 
-    model.forecast(forecast_full_ts)
+    if isinstance(model, get_args(ContextRequiredModelType)):
+        forecast_full_ts = ts.make_future(future_steps=5, tail_steps=model.context_size)
+        model.forecast(forecast_full_ts, prediction_size=5)
+    else:
+        forecast_full_ts = ts.make_future(future_steps=5)
+        model.forecast(forecast_full_ts)
 
     # forecasting only prefix
-    forecast_prefix_ts = ts.make_future(5)
-    forecast_prefix_ts.df = forecast_prefix_ts.df.iloc[:-2]
-
     torch.manual_seed(11)  # TODO: remove after fix at issue-802
-    model.forecast(forecast_prefix_ts)
+
+    if isinstance(model, get_args(ContextRequiredModelType)):
+        forecast_prefix_ts = ts.make_future(future_steps=5, tail_steps=model.context_size)
+        forecast_prefix_ts.df = forecast_prefix_ts.df.iloc[:-2]
+        model.forecast(forecast_prefix_ts, prediction_size=3)
+    else:
+        forecast_prefix_ts = ts.make_future(future_steps=5)
+        forecast_prefix_ts.df = forecast_prefix_ts.df.iloc[:-2]
+        model.forecast(forecast_prefix_ts)
 
     # checking
     forecast_full_df = forecast_full_ts.to_pandas()
@@ -97,13 +122,29 @@ def _test_forecast_out_sample_suffix(ts, model, transforms):
     model.fit(ts)
 
     # forecasting full
-    forecast_full_ts = ts.make_future(5)
-    model.forecast(forecast_full_ts)
+    if isinstance(model, get_args(ContextRequiredModelType)):
+        forecast_full_ts = ts.make_future(future_steps=5, tail_steps=model.context_size)
+        model.forecast(forecast_full_ts, prediction_size=5)
+    else:
+        forecast_full_ts = ts.make_future(future_steps=5)
+        model.forecast(forecast_full_ts)
 
     # forecasting only suffix
-    forecast_gap_ts = ts.make_future(5)
-    forecast_gap_ts.df = forecast_gap_ts.df.iloc[2:]
-    model.forecast(forecast_gap_ts)
+    if isinstance(model, get_args(ContextRequiredModelType)):
+        forecast_gap_ts = ts.make_future(future_steps=5, tail_steps=model.context_size)
+
+        # firstly we should forecast prefix to use it as a context
+        forecast_prefix_ts = deepcopy(forecast_gap_ts)
+        forecast_prefix_ts.df = forecast_prefix_ts.df.iloc[:-3]
+        model.forecast(forecast_prefix_ts, prediction_size=2)
+        forecast_gap_ts.df = forecast_gap_ts.df.combine_first(forecast_prefix_ts.df)
+
+        # forecast suffix with known context for it
+        model.forecast(forecast_gap_ts, prediction_size=3)
+    else:
+        forecast_gap_ts = ts.make_future(future_steps=5)
+        forecast_gap_ts.df = forecast_gap_ts.df.iloc[2:]
+        model.forecast(forecast_gap_ts)
 
     # checking
     forecast_full_df = forecast_full_ts.to_pandas()
@@ -111,31 +152,35 @@ def _test_forecast_out_sample_suffix(ts, model, transforms):
     assert_frame_equal(forecast_gap_df, forecast_full_df.iloc[2:])
 
 
-def _test_forecast_mixed_in_out_sample(ts, model, transforms):
+def _test_forecast_mixed_in_out_sample(ts, model, transforms, num_skip_points):
+    # skip context required model
+    if isinstance(model, get_args(ContextRequiredModelType)):
+        raise NotImplementedError("Context required model can't pass this test!")
+
     # fitting
     df = ts.to_pandas()
     ts.fit_transform(transforms)
     model.fit(ts)
 
     # forecasting mixed in-sample and out-sample
-    future_ts = ts.make_future(5)
+    future_ts = ts.make_future(future_steps=5)
     future_df = future_ts.to_pandas().loc[:, pd.IndexSlice[:, "target"]]
     df_full = pd.concat((df, future_df))
     forecast_full_ts = TSDataset(df=df_full, freq=future_ts.freq)
     forecast_full_ts.transform(ts.transforms)
     forecast_full_ts.df.loc[:, pd.IndexSlice[:, "target"]] = np.NaN
-    forecast_full_ts.df = forecast_full_ts.df.iloc[6:]
+    forecast_full_ts.df = forecast_full_ts.df.iloc[num_skip_points:]
     model.forecast(forecast_full_ts)
 
     # forecasting only in sample
     forecast_in_sample_ts = TSDataset(df, freq="D")
     forecast_in_sample_ts.transform(ts.transforms)
     forecast_in_sample_ts.df.loc[:, pd.IndexSlice[:, "target"]] = np.NaN
-    forecast_in_sample_ts.df = forecast_in_sample_ts.df.iloc[6:]
+    forecast_in_sample_ts.df = forecast_in_sample_ts.df.iloc[num_skip_points:]
     model.forecast(forecast_in_sample_ts)
 
     # forecasting only out sample
-    forecast_out_sample_ts = ts.make_future(5)
+    forecast_out_sample_ts = ts.make_future(future_steps=5)
     model.forecast(forecast_out_sample_ts)
 
     # checking
@@ -156,9 +201,6 @@ def _test_forecast_mixed_in_out_sample(ts, model, transforms):
         (HoltModel(), []),
         (HoltWintersModel(), []),
         (SimpleExpSmoothingModel(), []),
-        (MovingAverageModel(window=3), []),
-        (NaiveModel(lag=3), []),
-        (SeasonalMovingAverageModel(), []),
     ],
 )
 def test_forecast_in_sample_full(model, transforms, example_tsds):
@@ -180,6 +222,20 @@ def test_forecast_in_sample_full_failed(model, transforms, example_tsds):
     _test_forecast_in_sample_full(example_tsds, model, transforms)
 
 
+@pytest.mark.parametrize(
+    "model, transforms",
+    [
+        (MovingAverageModel(window=3), []),
+        (NaiveModel(lag=3), []),
+        (SeasonalMovingAverageModel(), []),
+        (DeadlineMovingAverageModel(window=1), []),
+    ],
+)
+def test_forecast_in_sample_full_failed_not_enough_context(model, transforms, example_tsds):
+    with pytest.raises(ValueError, match="Given context isn't big enough"):
+        _test_forecast_in_sample_full(example_tsds, model, transforms)
+
+
 @pytest.mark.parametrize(
     "model, transforms",
     [
@@ -236,10 +292,11 @@ def test_forecast_in_sample_full_not_implemented(model, transforms, example_tsds
         (MovingAverageModel(window=3), []),
         (NaiveModel(lag=3), []),
         (SeasonalMovingAverageModel(), []),
+        (DeadlineMovingAverageModel(window=1), []),
     ],
 )
 def test_forecast_in_sample_suffix(model, transforms, example_tsds):
-    _test_forecast_in_sample_suffix(example_tsds, model, transforms)
+    _test_forecast_in_sample_suffix(example_tsds, model, transforms, num_skip_points=50)
 
 
 @pytest.mark.parametrize(
@@ -277,7 +334,7 @@ def test_forecast_in_sample_suffix(model, transforms, example_tsds):
 )
 def test_forecast_in_sample_suffix_not_implemented(model, transforms, example_tsds):
     with pytest.raises(NotImplementedError, match="It is not possible to make in-sample predictions"):
-        _test_forecast_in_sample_suffix(example_tsds, model, transforms)
+        _test_forecast_in_sample_suffix(example_tsds, model, transforms, num_skip_points=50)
 
 
 @pytest.mark.parametrize(
@@ -298,6 +355,7 @@ def test_forecast_in_sample_suffix_not_implemented(model, transforms, example_ts
         (MovingAverageModel(window=3), []),
         (SeasonalMovingAverageModel(), []),
         (NaiveModel(lag=3), []),
+        (DeadlineMovingAverageModel(window=1), []),
         (BATSModel(use_trend=True), []),
         (TBATSModel(use_trend=True), []),
         (
@@ -349,6 +407,10 @@ def test_forecast_out_sample_prefix(model, transforms, example_tsds):
         (SimpleExpSmoothingModel(), []),
         (BATSModel(use_trend=True), []),
         (TBATSModel(use_trend=True), []),
+        (MovingAverageModel(window=3), []),
+        (SeasonalMovingAverageModel(), []),
+        (NaiveModel(lag=3), []),
+        (DeadlineMovingAverageModel(window=1), []),
     ],
 )
 def test_forecast_out_sample_suffix(model, transforms, example_tsds):
@@ -391,19 +453,6 @@ def test_forecast_out_sample_suffix_not_implemented(model, transforms, example_t
         _test_forecast_out_sample_suffix(example_tsds, model, transforms)
 
 
-@pytest.mark.xfail(strict=True)
-@pytest.mark.parametrize(
-    "model, transforms",
-    [
-        (MovingAverageModel(window=3), []),
-        (SeasonalMovingAverageModel(), []),
-        (NaiveModel(lag=3), []),
-    ],
-)
-def test_forecast_out_sample_suffix_failed(model, transforms, example_tsds):
-    _test_forecast_out_sample_suffix(example_tsds, model, transforms)
-
-
 @pytest.mark.parametrize(
     "model, transforms",
     [
@@ -422,7 +471,7 @@ def test_forecast_out_sample_suffix_failed(model, transforms, example_tsds):
     ],
 )
 def test_forecast_mixed_in_out_sample(model, transforms, example_tsds):
-    _test_forecast_mixed_in_out_sample(example_tsds, model, transforms)
+    _test_forecast_mixed_in_out_sample(example_tsds, model, transforms, num_skip_points=50)
 
 
 @pytest.mark.parametrize(
@@ -458,6 +507,19 @@ def test_forecast_mixed_in_out_sample(model, transforms, example_tsds):
         ),
     ],
 )
-def test_forecast_mixed_in_out_sample_not_implemented(model, transforms, example_tsds):
+def test_forecast_mixed_in_out_sample_not_implemented_in_sample(model, transforms, example_tsds):
     with pytest.raises(NotImplementedError, match="It is not possible to make in-sample predictions"):
-        _test_forecast_mixed_in_out_sample(example_tsds, model, transforms)
+        _test_forecast_mixed_in_out_sample(example_tsds, model, transforms, num_skip_points=50)
+
+
+@pytest.mark.parametrize(
+    "model, transforms",
+    [
+        (MovingAverageModel(window=3), []),
+        (SeasonalMovingAverageModel(), []),
+        (NaiveModel(lag=3), []),
+    ],
+)
+def test_forecast_mixed_in_out_sample_not_implemented_context_required(model, transforms, example_tsds):
+    with pytest.raises(NotImplementedError, match="Context required model can't pass this test"):
+        _test_forecast_mixed_in_out_sample(example_tsds, model, transforms, num_skip_points=50)

From 5b89c55e469edfcb3d495ac115b74e1671e6c391 Mon Sep 17 00:00:00 2001
From: "d.a.bunin" <d.a.bunin@tinkoff.ru>
Date: Mon, 5 Sep 2022 13:12:30 +0300
Subject: [PATCH 3/9] Fix PerSegmentModelMixin to clear values during in-sample

---
 etna/models/base.py | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/etna/models/base.py b/etna/models/base.py
index 2e0cc8c50..d204043a5 100644
--- a/etna/models/base.py
+++ b/etna/models/base.py
@@ -494,7 +494,9 @@ def _forecast(self, ts: TSDataset, **kwargs) -> TSDataset:
         result_df = result_df.set_index(["timestamp", "segment"])
         df = ts.to_pandas(flatten=True)
         df = df.set_index(["timestamp", "segment"])
-        # TODO: remember that it can be a trouble for in-sample forecasting
+        # clear values to be filled, otherwise during in-sample prediction our values won't be used
+        columns_to_clear = result_df.columns.intersection(df.columns)
+        df.loc[result_df.index, columns_to_clear] = np.NaN
         df = df.combine_first(result_df).reset_index()
 
         df = TSDataset.to_dataset(df)
@@ -508,7 +510,10 @@ def _forecast(self, ts: TSDataset, **kwargs) -> TSDataset:
 
 
 class MultiSegmentModelMixin(ModelForecastMixin):
-    """Mixin for holding methods for multi-segment prediction."""
+    """Mixin for holding methods for multi-segment prediction.
+
+    It currently isn't working with prediction intervals and context.
+    """
 
     def __init__(self, base_model: Any):
         """
@@ -557,7 +562,7 @@ def _forecast(self, ts: TSDataset, **kwargs) -> TSDataset:
         """
         horizon = len(ts.df)
         x = ts.to_pandas(flatten=True).drop(["segment"], axis=1)
-        # TODO: we haven't tested working with prediction_size here
+        # TODO: make it work with prediction intervals and context
         y = self._base_model.predict(x, **kwargs).reshape(-1, horizon).T
         ts.loc[:, pd.IndexSlice[:, "target"]] = y
         ts.inverse_transform()

From 05947811266e93ccb24da0e9f70d01b6251121c3 Mon Sep 17 00:00:00 2001
From: "d.a.bunin" <d.a.bunin@tinkoff.ru>
Date: Mon, 5 Sep 2022 15:39:07 +0300
Subject: [PATCH 4/9] Update changelog

---
 CHANGELOG.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index e2226c08d..92043b199 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -16,7 +16,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - 
 - 
 - 
-- 
+- Make `SeasonalMovingAverageModel` and `DeadlineMovingAverageModel` to work with context ([#917](https://github.com/tinkoff-ai/etna/pull/917))
 - 
 - 
 - 

From 0fd672534494fc8df4ef8d32540104975a8a7088 Mon Sep 17 00:00:00 2001
From: "d.a.bunin" <d.a.bunin@tinkoff.ru>
Date: Mon, 5 Sep 2022 15:47:56 +0300
Subject: [PATCH 5/9] Fix typo

---
 etna/models/deadline_ma.py | 2 +-
 etna/models/seasonal_ma.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/etna/models/deadline_ma.py b/etna/models/deadline_ma.py
index b193c51ab..ba17394e1 100644
--- a/etna/models/deadline_ma.py
+++ b/etna/models/deadline_ma.py
@@ -121,7 +121,7 @@ def predict(self, df: pd.DataFrame, prediction_size: int) -> np.ndarray:
         Raises
         ------
         ValueError:
-            if context isn't big enought
+            if context isn't big enough
         """
         context_beginning = self._get_context_beginning(df=df, prediction_size=prediction_size)
 
diff --git a/etna/models/seasonal_ma.py b/etna/models/seasonal_ma.py
index 87b08684d..97eb28c83 100644
--- a/etna/models/seasonal_ma.py
+++ b/etna/models/seasonal_ma.py
@@ -82,7 +82,7 @@ def predict(self, df: pd.DataFrame, prediction_size: int) -> np.ndarray:
         Raises
         ------
         ValueError:
-            if context isn't big enought
+            if context isn't big enough
         """
         expected_length = prediction_size + self.shift
         if len(df) < expected_length:

From 7eabe3da6c8b02543f0d6b050478abff55775a20 Mon Sep 17 00:00:00 2001
From: "d.a.bunin" <d.a.bunin@tinkoff.ru>
Date: Mon, 5 Sep 2022 16:38:24 +0300
Subject: [PATCH 6/9] Fix tests

---
 etna/analysis/outliers/prediction_interval_outliers.py     | 7 ++++---
 .../test_decomposition/test_stl_transform.py               | 4 ++--
 .../test_missing_values/test_impute_transform.py           | 6 +++---
 3 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/etna/analysis/outliers/prediction_interval_outliers.py b/etna/analysis/outliers/prediction_interval_outliers.py
index 839177093..c2921578f 100644
--- a/etna/analysis/outliers/prediction_interval_outliers.py
+++ b/etna/analysis/outliers/prediction_interval_outliers.py
@@ -86,9 +86,10 @@ def get_anomalies_prediction_interval(
         deepcopy(ts_inner), prediction_interval=True, quantiles=[lower_p, upper_p]
     )
     for segment in ts_inner.segments:
-        segment_slice = prediction_interval[:, segment, :][segment]
-        anomalies_mask = (segment_slice["target"] > segment_slice[f"target_{upper_p:.4g}"]) | (
-            segment_slice["target"] < segment_slice[f"target_{lower_p:.4g}"]
+        predicted_segment_slice = prediction_interval[:, segment, :][segment]
+        actual_segment_slice = ts_inner[:, segment, :][segment]
+        anomalies_mask = (actual_segment_slice["target"] > predicted_segment_slice[f"target_{upper_p:.4g}"]) | (
+            actual_segment_slice["target"] < predicted_segment_slice[f"target_{lower_p:.4g}"]
         )
         outliers_per_segment[segment] = list(time_points[anomalies_mask])
     return outliers_per_segment
diff --git a/tests/test_transforms/test_decomposition/test_stl_transform.py b/tests/test_transforms/test_decomposition/test_stl_transform.py
index 6f8988891..589b7823e 100644
--- a/tests/test_transforms/test_decomposition/test_stl_transform.py
+++ b/tests/test_transforms/test_decomposition/test_stl_transform.py
@@ -153,8 +153,8 @@ def test_forecast(ts_trend_seasonal, model_stl):
     ts_train.fit_transform(transforms=[transform])
     model = NaiveModel()
     model.fit(ts_train)
-    ts_future = ts_train.make_future(3)
-    ts_forecast = model.forecast(ts_future)
+    ts_future = ts_train.make_future(future_steps=3, tail_steps=model.context_size)
+    ts_forecast = model.forecast(ts_future, prediction_size=3)
     for segment in ts_forecast.segments:
         np.testing.assert_allclose(ts_forecast[:, segment, "target"], ts_test[:, segment, "target"], atol=0.1)
 
diff --git a/tests/test_transforms/test_missing_values/test_impute_transform.py b/tests/test_transforms/test_missing_values/test_impute_transform.py
index ea9154ff7..48da62bfc 100644
--- a/tests/test_transforms/test_missing_values/test_impute_transform.py
+++ b/tests/test_transforms/test_missing_values/test_impute_transform.py
@@ -343,9 +343,9 @@ def test_inverse_transform_in_forecast(df_with_missing_range_x_index_two_segment
     model = NaiveModel()
     ts.fit_transform(transforms=[imputer])
     model.fit(ts)
-    ts_test = ts.make_future(3)
-    assert np.all(ts_test[:, :, "target"].isna())
-    ts_forecast = model.forecast(ts_test)
+    ts_test = ts.make_future(future_steps=3, tail_steps=model.context_size)
+    assert np.all(ts_test[ts_test.index[-3] :, :, "target"].isna())
+    ts_forecast = model.forecast(ts_test, prediction_size=3)
     for segment in ts.segments:
         true_value = ts[:, segment, "target"].values[-1]
         assert np.all(ts_forecast[:, segment, "target"] == true_value)

From e8dbc86e426b648c1da680127a67fc56274eca28 Mon Sep 17 00:00:00 2001
From: "d.a.bunin" <d.a.bunin@tinkoff.ru>
Date: Mon, 5 Sep 2022 17:10:24 +0300
Subject: [PATCH 7/9] Clarify comment about filling with nans

---
 etna/models/base.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/etna/models/base.py b/etna/models/base.py
index d204043a5..40027f564 100644
--- a/etna/models/base.py
+++ b/etna/models/base.py
@@ -494,7 +494,7 @@ def _forecast(self, ts: TSDataset, **kwargs) -> TSDataset:
         result_df = result_df.set_index(["timestamp", "segment"])
         df = ts.to_pandas(flatten=True)
         df = df.set_index(["timestamp", "segment"])
-        # clear values to be filled, otherwise during in-sample prediction our values won't be used
+        # clear values to be filled, otherwise during in-sample prediction new values won't be set
         columns_to_clear = result_df.columns.intersection(df.columns)
         df.loc[result_df.index, columns_to_clear] = np.NaN
         df = df.combine_first(result_df).reset_index()

From d44a457e5b34f71807b10a82cda4ec382ca5aad1 Mon Sep 17 00:00:00 2001
From: "d.a.bunin" <d.a.bunin@tinkoff.ru>
Date: Tue, 6 Sep 2022 12:11:37 +0300
Subject: [PATCH 8/9] Fix comments on PR

---
 etna/models/deadline_ma.py              | 52 ++++++++++++++++++----
 tests/test_models/test_simple_models.py | 58 +++++++++++++++++++++++++
 2 files changed, 102 insertions(+), 8 deletions(-)

diff --git a/etna/models/deadline_ma.py b/etna/models/deadline_ma.py
index ba17394e1..8ec40fa50 100644
--- a/etna/models/deadline_ma.py
+++ b/etna/models/deadline_ma.py
@@ -83,18 +83,52 @@ def fit(self, df: pd.DataFrame, regressors: List[str]) -> "_DeadlineMovingAverag
 
         return self
 
-    def _get_context_beginning(self, df: pd.DataFrame, prediction_size: int):
+    @staticmethod
+    def _get_context_beginning(
+        df: pd.DataFrame, prediction_size: int, seasonality: SeasonalityMode, window: int
+    ) -> pd.Timestamp:
+        """
+        Get timestamp where context begins.
+
+        Parameters
+        ----------
+        df:
+            Time series in a long format.
+        prediction_size:
+            Number of last timestamps to leave after making prediction.
+            Previous timestamps will be used as a context for models that require it.
+        seasonality:
+            Seasonality.
+        window:
+            Number of values taken for forecast of each point.
+
+        Returns
+        -------
+        :
+            Timestamp with beginning of the context.
+
+        Raises
+        ------
+        ValueError:
+            if context isn't big enough
+        """
         df_history = df.iloc[:-prediction_size]
         history_timestamps = df_history["timestamp"]
         future_timestamps = df["timestamp"].iloc[-prediction_size:]
 
-        if self.seasonality == SeasonalityMode.month:
-            first_index = future_timestamps.iloc[0] - pd.DateOffset(months=self.window)
+        # if we have len(history_timestamps) == 0, then len(df) <= prediction_size
+        if len(history_timestamps) == 0:
+            raise ValueError(
+                "Given context isn't big enough, try to decrease context_size, prediction_size of increase length of given dataframe!"
+            )
+
+        if seasonality is SeasonalityMode.month:
+            first_index = future_timestamps.iloc[0] - pd.DateOffset(months=window)
 
-        elif self.seasonality == SeasonalityMode.year:
-            first_index = future_timestamps.iloc[0] - pd.DateOffset(years=self.window)
+        elif seasonality is SeasonalityMode.year:
+            first_index = future_timestamps.iloc[0] - pd.DateOffset(years=window)
 
-        if len(history_timestamps) == 0 or first_index < history_timestamps.iloc[0]:
+        if first_index < history_timestamps.iloc[0]:
             raise ValueError(
                 "Given context isn't big enough, try to decrease context_size, prediction_size of increase length of given dataframe!"
             )
@@ -123,14 +157,16 @@ def predict(self, df: pd.DataFrame, prediction_size: int) -> np.ndarray:
         ValueError:
             if context isn't big enough
         """
-        context_beginning = self._get_context_beginning(df=df, prediction_size=prediction_size)
+        context_beginning = self._get_context_beginning(
+            df=df, prediction_size=prediction_size, seasonality=self.seasonality, window=self.window
+        )
 
         df_history = df.iloc[:-prediction_size]
         history_targets = df_history["target"]
         history_timestamps = df_history["timestamp"]
-        future_timestamps = df["timestamp"].iloc[-prediction_size:]
         history_targets = history_targets.loc[history_timestamps >= context_beginning]
         history_timestamps = history_timestamps.loc[history_timestamps >= context_beginning]
+        future_timestamps = df["timestamp"].iloc[-prediction_size:]
 
         index = pd.date_range(start=context_beginning, end=future_timestamps.iloc[-1])
         res = np.append(history_targets.values, np.zeros(prediction_size))
diff --git a/tests/test_models/test_simple_models.py b/tests/test_models/test_simple_models.py
index 5fb9f77a7..c0a4e3725 100644
--- a/tests/test_models/test_simple_models.py
+++ b/tests/test_models/test_simple_models.py
@@ -5,6 +5,7 @@
 from etna.datasets import TSDataset
 from etna.metrics import MAE
 from etna.models.deadline_ma import DeadlineMovingAverageModel
+from etna.models.deadline_ma import SeasonalityMode
 from etna.models.deadline_ma import _DeadlineMovingAverageModel
 from etna.models.moving_average import MovingAverageModel
 from etna.models.naive import NaiveModel
@@ -54,6 +55,63 @@ def test_simple_model_forecaster_fail(simple_df):
         _ = sma_model.forecast(future_ts, prediction_size=7)
 
 
+@pytest.mark.parametrize(
+    "freq, periods, start, prediction_size, seasonality, window, expected",
+    [
+        ("D", 31 + 1, "2020-01-01", 1, SeasonalityMode.month, 1, pd.Timestamp("2020-01-01")),
+        ("D", 31 + 2, "2020-01-01", 1, SeasonalityMode.month, 1, pd.Timestamp("2020-01-02")),
+        ("D", 31 + 5, "2020-01-01", 5, SeasonalityMode.month, 1, pd.Timestamp("2020-01-01")),
+        ("D", 31 + 29 + 1, "2020-01-01", 1, SeasonalityMode.month, 2, pd.Timestamp("2020-01-01")),
+        ("D", 31 + 29 + 31 + 1, "2020-01-01", 1, SeasonalityMode.month, 3, pd.Timestamp("2020-01-01")),
+        ("H", 31 * 24 + 1, "2020-01-01", 1, SeasonalityMode.month, 1, pd.Timestamp("2020-01-01")),
+        ("H", 31 * 24 + 2, "2020-01-01", 1, SeasonalityMode.month, 1, pd.Timestamp("2020-01-01 01:00")),
+        ("H", 31 * 24 + 5, "2020-01-01", 5, SeasonalityMode.month, 1, pd.Timestamp("2020-01-01")),
+        ("H", (31 + 29) * 24 + 1, "2020-01-01", 1, SeasonalityMode.month, 2, pd.Timestamp("2020-01-01")),
+        ("H", (31 + 29 + 31) * 24 + 1, "2020-01-01", 1, SeasonalityMode.month, 3, pd.Timestamp("2020-01-01")),
+        ("D", 366 + 1, "2020-01-01", 1, SeasonalityMode.year, 1, pd.Timestamp("2020-01-01")),
+        ("D", 366 + 2, "2020-01-01", 1, SeasonalityMode.year, 1, pd.Timestamp("2020-01-02")),
+        ("D", 366 + 5, "2020-01-01", 5, SeasonalityMode.year, 1, pd.Timestamp("2020-01-01")),
+        ("D", 366 + 365 + 1, "2020-01-01", 1, SeasonalityMode.year, 2, pd.Timestamp("2020-01-01")),
+        ("D", 366 + 365 + 365 + 1, "2020-01-01", 1, SeasonalityMode.year, 3, pd.Timestamp("2020-01-01")),
+        ("H", 366 * 24 + 1, "2020-01-01", 1, SeasonalityMode.year, 1, pd.Timestamp("2020-01-01")),
+        ("H", 366 * 24 + 2, "2020-01-01", 1, SeasonalityMode.year, 1, pd.Timestamp("2020-01-01 01:00")),
+        ("H", 366 * 24 + 5, "2020-01-01", 5, SeasonalityMode.year, 1, pd.Timestamp("2020-01-01")),
+        ("H", (366 + 365) * 24 + 1, "2020-01-01", 1, SeasonalityMode.year, 2, pd.Timestamp("2020-01-01")),
+        ("H", (366 + 365 + 365) * 24 + 1, "2020-01-01", 1, SeasonalityMode.year, 3, pd.Timestamp("2020-01-01")),
+    ],
+)
+def test_deadline_get_context_beginning_ok(freq, periods, start, prediction_size, seasonality, window, expected):
+    df = pd.DataFrame({"timestamp": pd.date_range(start=start, periods=periods, freq=freq)})
+
+    obtained = _DeadlineMovingAverageModel._get_context_beginning(df, prediction_size, seasonality, window)
+
+    assert obtained == expected
+
+
+@pytest.mark.parametrize(
+    "freq, periods, start, prediction_size, seasonality, window",
+    [
+        ("D", 1, "2020-01-01", 1, SeasonalityMode.month, 1),
+        ("H", 1, "2020-01-01", 1, SeasonalityMode.month, 1),
+        ("D", 1, "2020-01-01", 1, SeasonalityMode.year, 1),
+        ("H", 1, "2020-01-01", 1, SeasonalityMode.year, 1),
+        ("D", 1, "2020-01-01", 2, SeasonalityMode.month, 1),
+        ("H", 1, "2020-01-01", 2, SeasonalityMode.month, 1),
+        ("D", 1, "2020-01-01", 2, SeasonalityMode.year, 1),
+        ("H", 1, "2020-01-01", 2, SeasonalityMode.year, 1),
+        ("D", 31 + 1, "2020-01-01", 2, SeasonalityMode.month, 1),
+        ("H", 31 * 24 + 1, "2020-01-01", 2, SeasonalityMode.month, 1),
+        ("D", 366 + 1, "2020-01-01", 2, SeasonalityMode.year, 1),
+        ("H", 366 * 24 + 1, "2020-01-01", 2, SeasonalityMode.year, 1),
+    ],
+)
+def test_deadline_get_context_beginning_fail(freq, periods, start, prediction_size, seasonality, window):
+    df = pd.DataFrame({"timestamp": pd.date_range(start=start, periods=periods, freq=freq)})
+
+    with pytest.raises(ValueError, match="Given context isn't big enough"):
+        _ = _DeadlineMovingAverageModel._get_context_beginning(df, prediction_size, seasonality, window)
+
+
 @pytest.mark.parametrize("model", [DeadlineMovingAverageModel])
 def test_deadline_model_forecaster_run(simple_df, model):
     model = model(window=1)

From ad6dc547bffe69c4c27c871e31beef138ed6043d Mon Sep 17 00:00:00 2001
From: "d.a.bunin" <d.a.bunin@tinkoff.ru>
Date: Tue, 6 Sep 2022 13:32:40 +0300
Subject: [PATCH 9/9] Fix test names

---
 tests/test_models/test_simple_models.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/tests/test_models/test_simple_models.py b/tests/test_models/test_simple_models.py
index c0a4e3725..cdfbf0e77 100644
--- a/tests/test_models/test_simple_models.py
+++ b/tests/test_models/test_simple_models.py
@@ -47,7 +47,7 @@ def test_simple_model_forecaster_run(simple_df, model):
     assert len(res) == 14
 
 
-def test_simple_model_forecaster_fail(simple_df):
+def test_simple_model_forecaster_fail_not_enough_context(simple_df):
     sma_model = SeasonalMovingAverageModel(window=1000, seasonality=7)
     sma_model.fit(simple_df)
     future_ts = simple_df.make_future(future_steps=7, tail_steps=sma_model.context_size)
@@ -105,7 +105,9 @@ def test_deadline_get_context_beginning_ok(freq, periods, start, prediction_size
         ("H", 366 * 24 + 1, "2020-01-01", 2, SeasonalityMode.year, 1),
     ],
 )
-def test_deadline_get_context_beginning_fail(freq, periods, start, prediction_size, seasonality, window):
+def test_deadline_get_context_beginning_fail_not_enough_context(
+    freq, periods, start, prediction_size, seasonality, window
+):
     df = pd.DataFrame({"timestamp": pd.date_range(start=start, periods=periods, freq=freq)})
 
     with pytest.raises(ValueError, match="Given context isn't big enough"):
@@ -123,7 +125,7 @@ def test_deadline_model_forecaster_run(simple_df, model):
     assert len(res) == 14
 
 
-def test_sdeadline_model_forecaster_fail(simple_df):
+def test_sdeadline_model_forecaster_fail_not_enough_context(simple_df):
     model = DeadlineMovingAverageModel(window=1000)
     model.fit(simple_df)
     future_ts = simple_df.make_future(future_steps=7, tail_steps=model.context_size)