Fix typing

tinkoff-ai · Mar 14, 2023 · 4f7231b · 4f7231b
1 parent 1ffa2e5
commit 4f7231b
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/test_models/test_mixins.py b/tests/test_models/test_mixins.py
@@ -27,7 +27,7 @@
 
 
 class DummyPredictAdapter(BaseAdapter):
-    def fit(self, df: pd.DataFrame, **kwargs) -> "DummyAdapter":
+    def fit(self, df: pd.DataFrame, **kwargs) -> "DummyPredictAdapter":
         return self
 
     def predict(self, df: pd.DataFrame, **kwargs) -> np.ndarray:
@@ -40,12 +40,12 @@ def predict_components(self, df: pd.DataFrame, **kwargs) -> pd.DataFrame:
         df = df.drop(columns=["target"])
         return df
 
-    def get_model(self) -> "DummyAdapter":
+    def get_model(self) -> "DummyPredictAdapter":
         return self
 
 
 class DummyForecastPredictAdapter(DummyPredictAdapter):
-    def fit(self, df: pd.DataFrame, **kwargs) -> "DummyAdapter":
+    def fit(self, df: pd.DataFrame, **kwargs) -> "DummyForecastPredictAdapter":
         return self
 
     def forecast(self, df: pd.DataFrame, **kwargs) -> np.ndarray: