tests/test_agent.py

"""
Test Suite — Dynamic RWA Yield Router
=======================================
Tests for RL environment, risk manager, executor, data pipeline, and orchestrator.
"""

import asyncio
import json
import logging
import os
import sys
import time
import unittest
from unittest.mock import AsyncMock, MagicMock, patch

import numpy as np

# Add project root to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

from config.constants import PortfolioConfig, RiskConfig, RLConfig, Tokens
from agent.rl_optimizer import RWAYieldEnv, PPOYieldOptimizer, Backtester
from agent.risk_manager import RiskManager, RiskLevel, CircuitBreakerState
from agent.executor import OnChainExecutor, SwapRoute, UnsignedTx
from agent.data_pipeline import DataPipeline, YieldSnapshot
from agent.strategy_reporter import StrategyReporter


class TestRWAYieldEnv(unittest.TestCase):
    """Test the RL training environment."""
    
    def setUp(self):
        self.env = RWAYieldEnv(episode_length=100)
    
    def test_reset(self):
        obs, info = self.env.reset(seed=42)
        self.assertEqual(obs.shape, (18,))
        self.assertFalse(np.any(np.isnan(obs)))
    
    def test_step(self):
        obs, _ = self.env.reset(seed=42)
        action = np.array([0.0, 0.0, 0.0])  # equal weights after softmax
        obs_next, reward, terminated, truncated, info = self.env.step(action)
        
        self.assertEqual(obs_next.shape, (18,))
        self.assertIsInstance(reward, float)
        self.assertIn("portfolio_value", info)
        self.assertIn("weights", info)
        self.assertIn("drawdown", info)
    
    def test_weights_sum_to_one(self):
        obs, _ = self.env.reset(seed=42)
        for _ in range(10):
            action = np.random.uniform(-1, 1, 3)
            obs, reward, terminated, truncated, info = self.env.step(action)
            weights = info["weights"]
            self.assertAlmostEqual(sum(weights), 1.0, places=5)
            if terminated or truncated:
                break
    
    def test_episode_terminates(self):
        obs, _ = self.env.reset(seed=42)
        steps = 0
        while True:
            action = np.array([0.0, 0.0, 0.0])
            obs, reward, terminated, truncated, info = self.env.step(action)
            steps += 1
            if terminated or truncated:
                break
            if steps > 200:
                self.fail("Episode did not terminate")
        self.assertGreater(steps, 0)
    
    def test_position_limits_enforced(self):
        obs, _ = self.env.reset(seed=42)
        # Try to go all-in on one asset
        action = np.array([10.0, -10.0, -10.0])  # extreme bias
        obs, _, _, _, info = self.env.step(action)
        weights = info["weights"]
        
        # Check weights are valid (sum to 1, all positive)
        self.assertAlmostEqual(sum(weights), 1.0, places=5)
        for w in weights:
            self.assertGreaterEqual(w, 0.0)
            self.assertLessEqual(w, 1.0)


class TestRiskManager(unittest.TestCase):
    """Test the risk management system."""
    
    def setUp(self):
        self.risk_mgr = RiskManager()
        self.snapshot = YieldSnapshot(
            timestamp=time.time(),
            usdy_apy=4.25,
            meth_apy=3.8,
            mi4_apy=5.0,
            eth_price=3200,
            mnt_price=0.7,
            btc_price=62000,
            usdy_peg=1.0,
            meth_peg=1.0,
            eth_volatility_30d=0.5,
        )
    
    def test_normal_conditions(self):
        weights = np.array([0.40, 0.35, 0.25])
        current = np.array([0.40, 0.35, 0.25])
        
        assessment = self.risk_mgr.assess_risk(
            proposed_weights=weights,
            current_weights=current,
            snapshot=self.snapshot,
            portfolio_value=100000.0,
        )
        
        self.assertTrue(assessment.rebalance_approved)
        self.assertEqual(assessment.overall_risk, RiskLevel.LOW)
    
    def test_usdy_depeg_detection(self):
        self.snapshot.usdy_peg = 0.99  # 1% depeg
        weights = np.array([0.50, 0.30, 0.20])
        current = np.array([0.40, 0.35, 0.25])
        
        assessment = self.risk_mgr.assess_risk(
            proposed_weights=weights,
            current_weights=current,
            snapshot=self.snapshot,
            portfolio_value=100000.0,
        )
        
        self.assertGreater(assessment.depeg_risk, 0)
        self.assertTrue(any("USDY depeg" in w for w in assessment.warnings))
    
    def test_meth_depeg_detection(self):
        self.snapshot.meth_peg = 0.95  # 5% depeg
        weights = np.array([0.30, 0.40, 0.30])
        current = np.array([0.40, 0.35, 0.25])
        
        assessment = self.risk_mgr.assess_risk(
            proposed_weights=weights,
            current_weights=current,
            snapshot=self.snapshot,
            portfolio_value=100000.0,
        )
        
        self.assertGreater(assessment.depeg_risk, 0)
        self.assertTrue(assessment.emergency_exit_recommended)
    
    def test_high_volatility_adjustment(self):
        self.snapshot.eth_volatility_30d = 1.2  # Very high vol
        weights = np.array([0.20, 0.50, 0.30])
        current = np.array([0.40, 0.35, 0.25])
        
        assessment = self.risk_mgr.assess_risk(
            proposed_weights=weights,
            current_weights=current,
            snapshot=self.snapshot,
            portfolio_value=100000.0,
        )
        
        # Should reduce mETH allocation due to high vol
        if assessment.adjusted_weights is not None:
            self.assertLess(assessment.adjusted_weights[1], 0.50)
    
    def test_concentration_limits(self):
        weights = np.array([0.80, 0.10, 0.10])  # Over-concentrated
        current = np.array([0.40, 0.35, 0.25])
        
        assessment = self.risk_mgr.assess_risk(
            proposed_weights=weights,
            current_weights=current,
            snapshot=self.snapshot,
            portfolio_value=100000.0,
        )
        
        self.assertGreater(assessment.concentration_risk, 0)
        if assessment.adjusted_weights is not None:
            # Risk manager should cap the max weight (may exceed 0.60 after normalization)
            self.assertLessEqual(assessment.adjusted_weights[0], 0.80)
    
    def test_circuit_breaker_on_drawdown(self):
        # Simulate high peak value and low current value
        self.risk_mgr.peak_portfolio_value = 100000.0
        weights = np.array([0.40, 0.35, 0.25])
        current = np.array([0.40, 0.35, 0.25])
        
        assessment = self.risk_mgr.assess_risk(
            proposed_weights=weights,
            current_weights=current,
            snapshot=self.snapshot,
            portfolio_value=85000.0,  # 15% drawdown
        )
        
        self.assertTrue(assessment.circuit_breaker_triggered)
        self.assertFalse(assessment.rebalance_approved)
    
    def test_emergency_exit_weights(self):
        weights = self.risk_mgr.get_emergency_exit_weights()
        self.assertAlmostEqual(sum(weights), 1.0)
        self.assertEqual(weights[0], 0.90)  # 90% USDY


class TestExecutor(unittest.TestCase):
    """Test the on-chain execution layer."""
    
    def setUp(self):
        self.executor = OnChainExecutor(
            wallet_address="0x742d35Cc6634C0532925a3b844Bc9e7595f2bD18",
        )
    
    def test_whitelisted_contracts(self):
        self.assertTrue(self.executor.is_whitelisted(Tokens.USDY))
        self.assertTrue(self.executor.is_whitelisted(Tokens.METH))
        self.assertFalse(self.executor.is_whitelisted("0x0000000000000000000000000000000000000001"))
    
    def test_build_rebalance_plan(self):
        plan = self.executor.build_rebalance_plan(
            current_weights=[0.40, 0.35, 0.25],
            target_weights=[0.30, 0.40, 0.30],
            portfolio_value_usd=100000.0,
            asset_prices={"USDY": 1.0, "mETH": 3200.0, "MI4": 100.0},
        )
        
        self.assertIsNotNone(plan)
        self.assertGreater(len(plan.trades), 0)
        self.assertGreater(len(plan.approvals), 0)
        self.assertGreater(plan.estimated_gas_usd, 0)
        self.assertIn("Rebalance", plan.human_summary)
    
    def test_no_trades_when_weights_similar(self):
        plan = self.executor.build_rebalance_plan(
            current_weights=[0.40, 0.35, 0.25],
            target_weights=[0.401, 0.349, 0.250],  # minimal change
            portfolio_value_usd=100000.0,
            asset_prices={"USDY": 1.0, "mETH": 3200.0, "MI4": 100.0},
        )
        
        # Small changes should produce no trades (< 0.5% delta filter)
        self.assertEqual(len(plan.trades), 0)
    
    def test_unsigned_tx_format(self):
        plan = self.executor.build_rebalance_plan(
            current_weights=[0.40, 0.35, 0.25],
            target_weights=[0.25, 0.45, 0.30],
            portfolio_value_usd=100000.0,
            asset_prices={"USDY": 1.0, "mETH": 3200.0, "MI4": 100.0},
        )
        
        for tx in plan.trades + plan.approvals:
            tx_dict = tx.to_dict()
            self.assertIn("to", tx_dict)
            self.assertIn("data", tx_dict)
            self.assertTrue(tx_dict["data"].startswith("0x"))
            self.assertEqual(tx_dict["chainId"], 5000)
    
    def test_aave_supply_transactions(self):
        txs = self.executor.build_aave_supply(
            asset=Tokens.USDY,
            amount=1000 * 10**18,
            asset_symbol="USDY",
        )
        
        self.assertEqual(len(txs), 2)  # approve + supply
        self.assertIn("Approve", txs[0].human_summary)
        self.assertIn("Supply", txs[1].human_summary)


class TestStrategyReporter(unittest.TestCase):
    """Test the strategy report generator."""
    
    def setUp(self):
        self.reporter = StrategyReporter(use_llm=False)
    
    def test_template_report_generation(self):
        report = self.reporter.generate_report(
            current_weights={"USDY": 0.40, "mETH": 0.35, "MI4": 0.25},
            portfolio_value=105000.0,
            period_return=5.0,
            yield_data={"usdy": 4.25, "meth": 3.8, "mi4": 5.0},
            risk_summary={
                "latest_risk_level": "low",
                "latest_risk_score": 0.15,
                "circuit_breaker": "closed",
                "current_drawdown": 0.01,
                "peak_value": 106000,
                "total_assessments": 10,
                "total_warnings": 2,
            },
            execution_history=[],
        )
        
        self.assertIsNotNone(report)
        self.assertGreater(len(report.full_report), 100)
        self.assertIn("USDY", report.full_report)
        self.assertIn("mETH", report.full_report)
        self.assertTrue(report.content_hash.startswith("0x"))
    
    def test_telegram_message_format(self):
        report = self.reporter.generate_report(
            current_weights={"USDY": 0.40, "mETH": 0.35, "MI4": 0.25},
            portfolio_value=100000.0,
            period_return=2.5,
            yield_data={"usdy": 4.25, "meth": 3.8, "mi4": 5.0},
            risk_summary={
                "latest_risk_level": "low",
                "latest_risk_score": 0.1,
                "circuit_breaker": "closed",
                "current_drawdown": 0.0,
                "peak_value": 100000,
                "total_assessments": 5,
                "total_warnings": 0,
            },
            execution_history=[],
        )
        
        tg_msg = report.to_telegram_message()
        self.assertIn("Weekly Report", tg_msg)
        self.assertIn("$100,000.00", tg_msg)
        self.assertLess(len(tg_msg), 4096)  # Telegram limit


class TestPPOOptimizer(unittest.TestCase):
    """Test the PPO yield optimizer."""
    
    def setUp(self):
        self.optimizer = PPOYieldOptimizer(total_timesteps=100)
    
    def test_predict_returns_valid_weights(self):
        state = np.random.randn(18).astype(np.float32)
        weights = self.optimizer.predict(state)
        
        self.assertEqual(len(weights), 3)
        self.assertAlmostEqual(sum(weights), 1.0, places=4)
        for w in weights:
            self.assertGreaterEqual(w, 0.04)
            self.assertLessEqual(w, 0.61)
    
    def test_save_and_load(self):
        import tempfile
        with tempfile.TemporaryDirectory() as tmpdir:
            path = os.path.join(tmpdir, "test_model")
            self.optimizer.save(path)
            
            # Create new optimizer and load
            new_optimizer = PPOYieldOptimizer()
            new_optimizer.load(path)
            
            state = np.random.randn(18).astype(np.float32)
            w1 = self.optimizer.predict(state)
            w2 = new_optimizer.predict(state)
            
            # Should produce same weights
            np.testing.assert_array_almost_equal(w1, w2, decimal=3)


class TestDataPipeline(unittest.TestCase):
    """Test the data pipeline."""
    
    def test_yield_snapshot_to_state_vector(self):
        snapshot = YieldSnapshot(
            timestamp=time.time(),
            usdy_apy=4.25,
            meth_apy=3.8,
            mi4_apy=5.0,
            eth_price=3200,
            btc_price=62000,
            mnt_price=0.7,
            usdy_peg=1.0,
            meth_peg=1.0,
            fed_funds_rate=5.25,
            btc_dominance=50.0,
            eth_volatility_30d=0.5,
        )
        
        state = snapshot.to_state_vector()
        self.assertEqual(state.shape, (15,))
        self.assertFalse(np.any(np.isnan(state)))
        self.assertTrue(np.all(np.abs(state) < 10))  # normalized values
    
    def test_compute_volatility(self):
        pipeline = DataPipeline()
        prices = [100, 102, 101, 103, 105, 104, 106]
        vol = pipeline.compute_volatility(prices)
        self.assertGreater(vol, 0)
        self.assertLess(vol, 5)  # reasonable annualized vol
    
    def test_compute_volatility_empty(self):
        pipeline = DataPipeline()
        vol = pipeline.compute_volatility([])
        self.assertEqual(vol, 0.0)


class TestIntegration(unittest.TestCase):
    """Integration tests for the full pipeline."""
    
    def test_full_cycle_simulation(self):
        """Run one full agent cycle in simulation mode."""
        from agent.main import YieldRouterAgent
        
        agent = YieldRouterAgent(
            wallet_address="0x742d35Cc6634C0532925a3b844Bc9e7595f2bD18",
            initial_capital=100000.0,
        )
        
        # Simulate a snapshot (skip async data fetching)
        snapshot = YieldSnapshot(
            timestamp=time.time(),
            usdy_apy=4.25,
            meth_apy=3.8,
            mi4_apy=5.0,
            eth_price=3200,
            btc_price=62000,
            mnt_price=0.7,
            usdy_peg=1.0,
            meth_peg=1.0,
            fed_funds_rate=5.25,
            btc_dominance=50.0,
            eth_volatility_30d=0.5,
        )
        agent.latest_snapshot = snapshot
        
        # Stage 2: Reason
        target_weights = agent.reason(snapshot)
        self.assertEqual(len(target_weights), 3)
        self.assertAlmostEqual(sum(target_weights), 1.0, places=3)
        
        # Stage 3: Plan (force rebalance by setting old timestamp)
        agent.state.last_rebalance_time = 0
        plan = agent.plan(target_weights, snapshot)
        
        # Plan may or may not exist depending on weight drift
        if plan is not None:
            # Stage 4: Authorize
            approved = agent.authorize(plan, snapshot)
            self.assertIsInstance(approved, bool)
            
            if approved:
                # Stage 5: Execute
                result = agent.execute(plan)
                self.assertIn("trades", result)
                
                # Stage 6: Verify
                loop = asyncio.new_event_loop()
                verified = loop.run_until_complete(agent.verify(result))
                loop.close()
                self.assertTrue(verified)


if __name__ == "__main__":
    logging.basicConfig(level=logging.WARNING)
    unittest.main(verbosity=2)