Spaces:
Paused
Paused
File size: 7,503 Bytes
1f2e77f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 | #!/usr/bin/env python3
"""
Phase 2 Rate Limiting Test - Market Data Hub & Global Rate Budget
Tests the new Phase 2 optimizations:
1. MarketDataHub prefetching and request coalescing
2. GlobalRateBudget priority allocation
3. Integration with AlpacaClient
"""
from aureon.core.aureon_baton_link import link_system as _baton_link; _baton_link(__name__)
import time
import logging
import sys
from typing import List
# Setup logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
def test_market_data_hub():
"""Test MarketDataHub functionality."""
logger.info("Testing MarketDataHub...")
try:
from aureon.exchanges.alpaca_client import AlpacaClient
client = AlpacaClient()
# Start market data hub
client.start_market_data_hub()
logger.info("MarketDataHub started")
# Test symbols
test_symbols = ['BTC/USD', 'ETH/USD', 'SOL/USD']
# First requests should hit API
logger.info("First round of quote requests (should hit API)...")
start_time = time.time()
quotes1 = {}
for symbol in test_symbols:
quotes1[symbol] = client.get_last_quote(symbol)
logger.info(f" {symbol}: {quotes1[symbol].get('last', {}).get('price', 'N/A')}")
first_round_time = time.time() - start_time
logger.info(f"First round took {first_round_time:.2f}s")
# Wait for prefetch cache to populate
time.sleep(2.5)
# Second requests should hit cache
logger.info("Second round of quote requests (should hit cache)...")
start_time = time.time()
quotes2 = {}
for symbol in test_symbols:
quotes2[symbol] = client.get_last_quote(symbol)
logger.info(f" {symbol}: {quotes2[symbol].get('last', {}).get('price', 'N/A')}")
second_round_time = time.time() - start_time
logger.info(f"Second round took {second_round_time:.2f}s")
# Check if cache improved performance
if second_round_time < first_round_time * 0.5:
logger.info("✅ Cache performance improvement detected!")
else:
logger.warning("⚠️ Cache performance not significantly improved")
# Get hub stats
stats = client.get_market_data_hub_stats()
logger.info(f"MarketDataHub stats: {stats}")
# Stop hub
client.stop_market_data_hub()
logger.info("MarketDataHub stopped")
return True
except Exception as e:
logger.error(f"MarketDataHub test failed: {e}")
return False
def test_request_coalescing():
"""Test request coalescing functionality."""
logger.info("Testing request coalescing...")
try:
import threading
from aureon.exchanges.alpaca_client import AlpacaClient
client = AlpacaClient()
client.start_market_data_hub()
results = {}
errors = []
def fetch_quote(symbol: str, result_dict: dict, error_list: list):
try:
quote = client.get_last_quote(symbol)
result_dict[symbol] = quote.get('last', {}).get('price')
except Exception as e:
error_list.append(f"{symbol}: {e}")
# Test coalescing with near-simultaneous requests
test_symbol = 'BTC/USD'
threads = []
logger.info(f"Starting 5 simultaneous requests for {test_symbol}...")
start_time = time.time()
for i in range(5):
thread = threading.Thread(
target=fetch_quote,
args=(test_symbol, results, errors)
)
threads.append(thread)
# Start all threads within 10ms
for thread in threads:
thread.start()
time.sleep(0.002) # 2ms stagger
# Wait for all to complete
for thread in threads:
thread.join()
elapsed = time.time() - start_time
logger.info(f"All requests completed in {elapsed:.2f}s")
logger.info(f"Results: {len(results)} successful, {len(errors)} errors")
if len(results) >= 4 and elapsed < 0.5: # Should complete quickly due to coalescing
logger.info("✅ Request coalescing working!")
else:
logger.warning("⚠️ Request coalescing may not be working optimally")
client.stop_market_data_hub()
return len(errors) == 0
except Exception as e:
logger.error(f"Request coalescing test failed: {e}")
return False
def test_global_rate_budget():
"""Test GlobalRateBudget priority allocation."""
logger.info("Testing GlobalRateBudget...")
try:
from aureon.exchanges.alpaca_client import AlpacaClient
client = AlpacaClient()
# Get initial stats
stats_before = client.get_global_rate_budget_stats()
logger.info(f"Initial GlobalRateBudget stats: {stats_before}")
# Make some requests to generate stats
logger.info("Making test requests...")
# Quotes (lowest priority)
client.get_last_quote('BTC/USD')
client.get_last_quote('ETH/USD')
# Positions (medium priority)
client.get_positions()
# Account (medium priority)
client.get_account()
# Get stats after
stats_after = client.get_global_rate_budget_stats()
logger.info(f"Final GlobalRateBudget stats: {stats_after}")
# Check if stats were updated
if stats_after != stats_before:
logger.info("✅ GlobalRateBudget tracking requests!")
else:
logger.warning("⚠️ GlobalRateBudget stats not updating")
return True
except Exception as e:
logger.error(f"GlobalRateBudget test failed: {e}")
return False
def main():
"""Run all Phase 2 tests."""
logger.info("Starting Phase 2 Rate Limiting Tests")
logger.info("=" * 50)
tests = [
("MarketDataHub", test_market_data_hub),
("Request Coalescing", test_request_coalescing),
("GlobalRateBudget", test_global_rate_budget),
]
results = []
for test_name, test_func in tests:
logger.info(f"\n🧪 Running {test_name} test...")
try:
success = test_func()
results.append((test_name, success))
status = "✅ PASSED" if success else "❌ FAILED"
logger.info(f"{test_name}: {status}")
except Exception as e:
logger.error(f"{test_name}: ❌ FAILED with exception: {e}")
results.append((test_name, False))
logger.info("\n" + "=" * 50)
logger.info("Phase 2 Test Results Summary:")
all_passed = True
for test_name, success in results:
status = "✅ PASSED" if success else "❌ FAILED"
logger.info(f" {test_name}: {status}")
if not success:
all_passed = False
if all_passed:
logger.info("\n🎉 All Phase 2 optimizations are working!")
logger.info("🚀 Ready to reduce 429 errors significantly.")
else:
logger.warning("\n⚠️ Some Phase 2 optimizations may need attention.")
return 0 if all_passed else 1
if __name__ == "__main__":
sys.exit(main()) |