404 lines
14 KiB
Python
404 lines
14 KiB
Python
"""
|
|
Tests for the strategy generation pipeline.
|
|
Tests the flow: AI description → Blockly XML → JSON → Python
|
|
"""
|
|
import json
|
|
import pytest
|
|
import subprocess
|
|
import xml.etree.ElementTree as ET
|
|
import sys
|
|
import os
|
|
|
|
# Add src to path for imports
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))
|
|
|
|
from PythonGenerator import PythonGenerator
|
|
|
|
|
|
class TestStrategyBuilder:
|
|
"""Tests for the CmdForge strategy-builder tool."""
|
|
|
|
@staticmethod
|
|
def _get_blocks(root):
|
|
"""Get all block elements, handling XML namespaces."""
|
|
# Blockly uses namespace https://developers.google.com/blockly/xml
|
|
# Elements may be prefixed with {namespace}block or just block
|
|
blocks = []
|
|
for elem in root.iter():
|
|
# Get local name without namespace
|
|
tag = elem.tag.split('}')[-1] if '}' in elem.tag else elem.tag
|
|
if tag == 'block':
|
|
blocks.append(elem)
|
|
return blocks
|
|
|
|
@pytest.mark.integration
|
|
def test_simple_rsi_strategy(self):
|
|
"""Test generating a simple RSI-based strategy."""
|
|
input_data = {
|
|
"description": "Buy when RSI is below 30 and sell when RSI is above 70",
|
|
"indicators": [{"name": "RSI", "outputs": ["RSI"]}],
|
|
"signals": [],
|
|
"default_source": {"exchange": "binance", "market": "BTC/USDT", "timeframe": "5m"}
|
|
}
|
|
|
|
result = subprocess.run(
|
|
['strategy-builder'],
|
|
input=json.dumps(input_data),
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=120
|
|
)
|
|
|
|
assert result.returncode == 0, f"Strategy builder failed: {result.stderr}"
|
|
|
|
# Validate XML structure
|
|
xml_output = result.stdout.strip()
|
|
root = ET.fromstring(xml_output)
|
|
|
|
# Check it has the required structure
|
|
root_tag = root.tag.split('}')[-1] if '}' in root.tag else root.tag
|
|
assert root_tag == 'xml', f"Root should be 'xml', got {root_tag}"
|
|
|
|
blocks = self._get_blocks(root)
|
|
assert len(blocks) >= 2, f"Should have at least 2 blocks (buy and sell), got {len(blocks)}"
|
|
|
|
# Check for execute_if blocks
|
|
execute_if_blocks = [b for b in blocks if b.get('type') == 'execute_if']
|
|
assert len(execute_if_blocks) >= 2, "Should have at least 2 execute_if blocks"
|
|
|
|
# Check for trade_action blocks
|
|
trade_action_blocks = [b for b in blocks if b.get('type') == 'trade_action']
|
|
assert len(trade_action_blocks) >= 2, "Should have buy and sell trade actions"
|
|
|
|
# Check for indicator blocks
|
|
indicator_blocks = [b for b in blocks if 'indicator_RSI' in (b.get('type') or '')]
|
|
assert len(indicator_blocks) >= 2, "Should use RSI indicator"
|
|
|
|
@pytest.mark.integration
|
|
def test_ema_crossover_strategy(self):
|
|
"""Test generating an EMA crossover strategy."""
|
|
input_data = {
|
|
"description": "Buy when EMA 20 crosses above EMA 50, sell when EMA 20 crosses below EMA 50",
|
|
"indicators": [
|
|
{"name": "EMA_20", "outputs": ["ema"]},
|
|
{"name": "EMA_50", "outputs": ["ema"]}
|
|
],
|
|
"signals": [],
|
|
"default_source": {"exchange": "binance", "market": "ETH/USDT", "timeframe": "1h"}
|
|
}
|
|
|
|
result = subprocess.run(
|
|
['strategy-builder'],
|
|
input=json.dumps(input_data),
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=120
|
|
)
|
|
|
|
assert result.returncode == 0, f"Strategy builder failed: {result.stderr}"
|
|
|
|
xml_output = result.stdout.strip()
|
|
root = ET.fromstring(xml_output)
|
|
|
|
# Check for EMA indicator blocks
|
|
blocks = self._get_blocks(root)
|
|
ema_20_blocks = [b for b in blocks if 'indicator_EMA_20' in (b.get('type') or '')]
|
|
ema_50_blocks = [b for b in blocks if 'indicator_EMA_50' in (b.get('type') or '')]
|
|
|
|
assert len(ema_20_blocks) >= 1, "Should use EMA_20 indicator"
|
|
assert len(ema_50_blocks) >= 1, "Should use EMA_50 indicator"
|
|
|
|
@pytest.mark.integration
|
|
def test_no_indicators_price_only(self):
|
|
"""Test generating a price-based strategy without indicators."""
|
|
input_data = {
|
|
"description": "Buy when price drops 5% from previous candle, sell when price rises 3%",
|
|
"indicators": [],
|
|
"signals": [],
|
|
"default_source": {"exchange": "binance", "market": "BTC/USDT", "timeframe": "5m"}
|
|
}
|
|
|
|
result = subprocess.run(
|
|
['strategy-builder'],
|
|
input=json.dumps(input_data),
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=120
|
|
)
|
|
|
|
assert result.returncode == 0, f"Strategy builder failed: {result.stderr}"
|
|
|
|
xml_output = result.stdout.strip()
|
|
root = ET.fromstring(xml_output)
|
|
|
|
# Should be valid XML
|
|
assert root is not None
|
|
|
|
@pytest.mark.integration
|
|
def test_missing_indicator_error(self):
|
|
"""Test that strategy mentioning indicators without config fails."""
|
|
input_data = {
|
|
"description": "Buy when RSI is below 30",
|
|
"indicators": [], # No indicators configured
|
|
"signals": [],
|
|
"default_source": {"exchange": "binance", "market": "BTC/USDT", "timeframe": "5m"}
|
|
}
|
|
|
|
result = subprocess.run(
|
|
['strategy-builder'],
|
|
input=json.dumps(input_data),
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=120
|
|
)
|
|
|
|
assert result.returncode != 0, "Should fail when indicators mentioned but not configured"
|
|
assert "indicator" in result.stderr.lower(), "Error should mention indicators"
|
|
|
|
|
|
class TestPythonGenerator:
|
|
"""Tests for the PythonGenerator class."""
|
|
|
|
def test_simple_execute_if(self):
|
|
"""Test generating Python from a simple execute_if block."""
|
|
strategy_json = {
|
|
"type": "strategy",
|
|
"statements": [
|
|
{
|
|
"type": "execute_if",
|
|
"inputs": {
|
|
"CONDITION": {
|
|
"type": "comparison",
|
|
"operator": ">",
|
|
"inputs": {
|
|
"LEFT": {"type": "current_price"},
|
|
"RIGHT": {"type": "dynamic_value", "values": [50000]}
|
|
}
|
|
}
|
|
},
|
|
"statements": {
|
|
"DO": [
|
|
{
|
|
"type": "trade_action",
|
|
"trade_type": "buy",
|
|
"inputs": {"size": 0.01}
|
|
}
|
|
]
|
|
}
|
|
}
|
|
]
|
|
}
|
|
|
|
generator = PythonGenerator(
|
|
default_source={"exchange": "binance", "market": "BTC/USDT", "timeframe": "5m"},
|
|
strategy_id="test_strategy"
|
|
)
|
|
|
|
result = generator.generate(strategy_json)
|
|
|
|
assert "generated_code" in result
|
|
code = result["generated_code"]
|
|
|
|
# Check for expected code elements
|
|
assert "def next():" in code
|
|
assert "if " in code
|
|
assert "get_current_price" in code
|
|
assert "trade_order" in code
|
|
|
|
def test_indicator_condition(self):
|
|
"""Test generating Python with indicator conditions."""
|
|
strategy_json = {
|
|
"type": "strategy",
|
|
"statements": [
|
|
{
|
|
"type": "execute_if",
|
|
"inputs": {
|
|
"CONDITION": {
|
|
"type": "comparison",
|
|
"operator": "<",
|
|
"inputs": {
|
|
"LEFT": {
|
|
"type": "indicator_RSI",
|
|
"fields": {"OUTPUT": "RSI"}
|
|
},
|
|
"RIGHT": {"type": "dynamic_value", "values": [30]}
|
|
}
|
|
}
|
|
},
|
|
"statements": {
|
|
"DO": [
|
|
{
|
|
"type": "trade_action",
|
|
"trade_type": "buy",
|
|
"inputs": {"size": 0.1}
|
|
}
|
|
]
|
|
}
|
|
}
|
|
]
|
|
}
|
|
|
|
generator = PythonGenerator(
|
|
default_source={"exchange": "binance", "market": "BTC/USDT", "timeframe": "5m"},
|
|
strategy_id="test_indicator"
|
|
)
|
|
|
|
result = generator.generate(strategy_json)
|
|
code = result["generated_code"]
|
|
|
|
# Check for indicator processing
|
|
assert "process_indicator" in code
|
|
assert "RSI" in code
|
|
|
|
# Check indicators are tracked
|
|
assert len(result["indicators"]) > 0
|
|
|
|
def test_logical_and_condition(self):
|
|
"""Test generating Python with logical AND conditions."""
|
|
strategy_json = {
|
|
"type": "strategy",
|
|
"statements": [
|
|
{
|
|
"type": "execute_if",
|
|
"inputs": {
|
|
"CONDITION": {
|
|
"type": "logical_and",
|
|
"inputs": {
|
|
"left": {
|
|
"type": "comparison",
|
|
"operator": "<",
|
|
"inputs": {
|
|
"LEFT": {"type": "indicator_RSI", "fields": {"OUTPUT": "RSI"}},
|
|
"RIGHT": {"type": "dynamic_value", "values": [30]}
|
|
}
|
|
},
|
|
"right": {
|
|
"type": "flag_is_set",
|
|
"flag_name": "bought",
|
|
"flag_value": False
|
|
}
|
|
}
|
|
}
|
|
},
|
|
"statements": {
|
|
"DO": [
|
|
{"type": "trade_action", "trade_type": "buy", "inputs": {"size": 0.01}}
|
|
]
|
|
}
|
|
}
|
|
]
|
|
}
|
|
|
|
generator = PythonGenerator(
|
|
default_source={"exchange": "binance", "market": "BTC/USDT", "timeframe": "5m"},
|
|
strategy_id="test_and"
|
|
)
|
|
|
|
result = generator.generate(strategy_json)
|
|
code = result["generated_code"]
|
|
|
|
# Check for logical AND
|
|
assert " and " in code
|
|
assert "flags.get" in code
|
|
|
|
def test_set_flag(self):
|
|
"""Test generating Python for flag setting."""
|
|
strategy_json = {
|
|
"type": "strategy",
|
|
"statements": [
|
|
{
|
|
"type": "set_flag",
|
|
"flag_name": "bought",
|
|
"flag_value": "True"
|
|
}
|
|
]
|
|
}
|
|
|
|
generator = PythonGenerator(
|
|
default_source={"exchange": "binance", "market": "BTC/USDT", "timeframe": "5m"},
|
|
strategy_id="test_flag"
|
|
)
|
|
|
|
result = generator.generate(strategy_json)
|
|
code = result["generated_code"]
|
|
|
|
# Check for flag setting
|
|
assert "flags['bought']" in code
|
|
assert "True" in code
|
|
|
|
# Check flag is tracked
|
|
assert "bought" in result["flags_used"]
|
|
|
|
def test_trade_action_with_options(self):
|
|
"""Test generating Python for trade with stop loss and take profit."""
|
|
strategy_json = {
|
|
"type": "strategy",
|
|
"statements": [
|
|
{
|
|
"type": "trade_action",
|
|
"trade_type": "buy",
|
|
"inputs": {"size": 0.1},
|
|
"trade_options": [
|
|
{"type": "stop_loss", "inputs": {"stop_loss": 45000}},
|
|
{"type": "take_profit", "inputs": {"take_profit": 55000}}
|
|
]
|
|
}
|
|
]
|
|
}
|
|
|
|
generator = PythonGenerator(
|
|
default_source={"exchange": "binance", "market": "BTC/USDT", "timeframe": "5m"},
|
|
strategy_id="test_trade_options"
|
|
)
|
|
|
|
result = generator.generate(strategy_json)
|
|
code = result["generated_code"]
|
|
|
|
# Check for trade order call
|
|
assert "trade_order" in code
|
|
assert "buy" in code
|
|
|
|
def test_math_operation(self):
|
|
"""Test generating Python for math operations."""
|
|
strategy_json = {
|
|
"type": "strategy",
|
|
"statements": [
|
|
{
|
|
"type": "execute_if",
|
|
"inputs": {
|
|
"CONDITION": {
|
|
"type": "comparison",
|
|
"operator": ">",
|
|
"inputs": {
|
|
"LEFT": {"type": "current_price"},
|
|
"RIGHT": {
|
|
"type": "math_operation",
|
|
"inputs": {
|
|
"operator": "MULTIPLY",
|
|
"left_operand": {"type": "indicator_SMA", "fields": {"OUTPUT": "sma"}},
|
|
"right_operand": {"type": "dynamic_value", "values": [1.02]}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
},
|
|
"statements": {
|
|
"DO": [
|
|
{"type": "trade_action", "trade_type": "sell", "inputs": {"size": 0.01}}
|
|
]
|
|
}
|
|
}
|
|
]
|
|
}
|
|
|
|
generator = PythonGenerator(
|
|
default_source={"exchange": "binance", "market": "BTC/USDT", "timeframe": "5m"},
|
|
strategy_id="test_math"
|
|
)
|
|
|
|
result = generator.generate(strategy_json)
|
|
code = result["generated_code"]
|
|
|
|
# Check for math operation
|
|
assert "*" in code # Multiply operator
|