From 66f7c913570b29be29d180caa1e50ff75f2438b2 Mon Sep 17 00:00:00 2001
From: hippocritical <gunngirspam@gmail.com>
Date: Sun, 1 Jan 2023 22:03:45 +0100
Subject: [PATCH] Adding tests added more code inside NameUpdater to grab more
 variables.

---
 freqtrade/strategy/strategyupdater.py | 18 +++++++++-
 requirements.txt                      | 20 +++++++++++
 tests/test_strategy_updater.py        | 50 +++++++++++++++++++++------
 3 files changed, 77 insertions(+), 11 deletions(-)

diff --git a/freqtrade/strategy/strategyupdater.py b/freqtrade/strategy/strategyupdater.py
index 5b4bb8be0..1a0423076 100644
--- a/freqtrade/strategy/strategyupdater.py
+++ b/freqtrade/strategy/strategyupdater.py
@@ -107,12 +107,16 @@ class NameUpdater(ast.NodeTransformer):
             for field_name, field_value in ast.iter_fields(node):
                 self.check_strategy_and_config_settings(node, field_value)
                 self.check_fields(field_value)
+                for child in ast.iter_child_nodes(node):
+                    self.generic_visit(child)
 
     def check_fields(self, field_value):
         if isinstance(field_value, list):
             for item in field_value:
-                if isinstance(item, ast.AST):
+                if isinstance(item, ast.AST) or isinstance(item, ast.If):
                     self.visit(item)
+        if isinstance(field_value, ast.Name):
+            self.visit_Name(field_value)
 
     def check_strategy_and_config_settings(self, node, field_value):
         if (isinstance(field_value, ast.AST) and
@@ -157,6 +161,11 @@ class NameUpdater(ast.NodeTransformer):
         #        node.module = "freqtrade.strategy"
         return node
 
+    def visit_If(self, node: ast.If):
+        for child in ast.iter_child_nodes(node):
+            self.visit(child)
+        return self.generic_visit(node)
+
     def visit_FunctionDef(self, node):
         # if the function name is in the mapping, update it
         if node.name in StrategyUpdater.function_mapping:
@@ -165,6 +174,13 @@ class NameUpdater(ast.NodeTransformer):
             self.check_args(node)
         return self.generic_visit(node)
 
+    def visit_Assign(self, node):
+        if hasattr(node, "targets") and isinstance(node.targets, list):
+            for target in node.targets:
+                if hasattr(target, "id") and target.id in StrategyUpdater.name_mapping:
+                    target.id = StrategyUpdater.name_mapping[target.id]
+        return node
+
     def visit_Attribute(self, node):
         # if the attribute name is 'nr_of_successful_buys',
         # update it to 'nr_of_successful_entries'
diff --git a/requirements.txt b/requirements.txt
index 90bc4f702..3229ec3c4 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -57,3 +57,23 @@ schedule==1.1.0
 #WS Messages
 websockets==10.4
 janus==1.0.0
+
+pytest~=7.2.0
+freqtrade~=2022.12.dev0
+filelock~=3.8.2
+plotly~=5.11.0
+setuptools~=65.6.3
+starlette~=0.22.0
+gym~=0.21.0
+torch~=1.13.1
+scikit-learn~=1.1.3
+scipy~=1.9.3
+xgboost~=1.7.2
+catboost~=1.1.1
+lightgbm~=3.3.3
+astor~=0.8.1
+ta~=0.10.2
+finta~=1.3
+tapy~=1.9.1
+matplotlib~=3.6.2
+PyYAML~=6.0
diff --git a/tests/test_strategy_updater.py b/tests/test_strategy_updater.py
index cf18fcc25..5736e5c76 100644
--- a/tests/test_strategy_updater.py
+++ b/tests/test_strategy_updater.py
@@ -4,11 +4,11 @@ from freqtrade.strategy.strategyupdater import StrategyUpdater
 
 
 def test_strategy_updater(default_conf, caplog) -> None:
-    modified_code5 = StrategyUpdater.update_code(StrategyUpdater, """
-def confirm_trade_exit(sell_reason: str):
-    pass
+    modified_code2 = StrategyUpdater.update_code(StrategyUpdater, """
+ticker_interval = '15m'
+buy_some_parameter = IntParameter(space='buy')
+sell_some_parameter = IntParameter(space='sell')
 """)
-
     modified_code1 = StrategyUpdater.update_code(StrategyUpdater, """
 class testClass(IStrategy):
     def populate_buy_trend():
@@ -21,12 +21,6 @@ class testClass(IStrategy):
         pass
     def custom_sell():
         pass
-""")
-
-    modified_code2 = StrategyUpdater.update_code(StrategyUpdater, """
-buy_some_parameter = IntParameter(space='buy')
-sell_some_parameter = IntParameter(space='sell')
-ticker_interval = '15m'
 """)
     modified_code3 = StrategyUpdater.update_code(StrategyUpdater, """
 use_sell_signal = True
@@ -38,6 +32,32 @@ forcebuy_enable = True
     modified_code4 = StrategyUpdater.update_code(StrategyUpdater, """
 dataframe.loc[reduce(lambda x, y: x & y, conditions), ["buy", "buy_tag"]] = (1, "buy_signal_1")
 dataframe.loc[reduce(lambda x, y: x & y, conditions), 'sell'] = 1
+""")
+    modified_code5 = StrategyUpdater.update_code(StrategyUpdater, """
+def confirm_trade_exit(sell_reason: str):
+    pass
+    """)
+    modified_code6 = StrategyUpdater.update_code(StrategyUpdater, """
+order_time_in_force = {
+    'buy': 'gtc',
+    'sell': 'ioc'
+}
+order_types = {
+    'buy': 'limit',
+    'sell': 'market',
+    'stoploss': 'market',
+    'stoploss_on_exchange': False
+}
+unfilledtimeout = {
+    'buy': 1,
+    'sell': 2
+}
+""")
+
+    modified_code7 = StrategyUpdater.update_code(StrategyUpdater, """
+def confirm_trade_exit(sell_reason):
+    if (sell_reason == 'stop_loss'):
+        pass
 """)
 
     assert "populate_entry_trend" in modified_code1
@@ -63,3 +83,13 @@ dataframe.loc[reduce(lambda x, y: x & y, conditions), 'sell'] = 1
     assert "enter_tag" in modified_code4
 
     assert "exit_reason" in modified_code5
+
+    assert "'entry': 'gtc'" in modified_code6
+    assert "'exit': 'ioc'" in modified_code6
+    assert "'entry': 'limit'" in modified_code6
+    assert "'exit': 'market'" in modified_code6
+    assert "'entry': 1" in modified_code6
+    assert "'exit': 2" in modified_code6
+
+    assert "exit_reason" in modified_code7
+    assert "exit_reason == 'stop_loss'" in modified_code7