"""
Tests for requests library interception.
"""

import unittest
from unittest.mock import patch, MagicMock
import requests
from kindred_ai import Interceptor


class TestRequestsInterceptor(unittest.TestCase):
    """Test cases for requests library interception."""
    
    def setUp(self):
        """Set up test fixtures."""
        # Stop any existing interceptor
        if Interceptor.is_active():
            Interceptor.stop()
    
    def tearDown(self):
        """Clean up after tests."""
        if Interceptor.is_active():
            Interceptor.stop()
    
    def test_url_rewriting(self):
        """Test that URLs are correctly rewritten."""
        rules = {
            "api.openai.com": "http://localhost:8080/openai"
        }
        
        Interceptor.init(routing_rules=rules, verbose=False)
        
        # Mock at the adapter level to avoid replacing our wrapper
        with patch('requests.adapters.HTTPAdapter.send') as mock_send:
            from requests.models import Response
            mock_response = Response()
            mock_response.status_code = 200
            mock_response.reason = 'OK'
            mock_response.headers = {}
            mock_response.url = "http://localhost:8080/openai/v1/models"
            mock_response.raw = MagicMock()
            mock_response.raw.release_conn = MagicMock()
            mock_send.return_value = mock_response
            
            session = requests.Session()
            session.get("https://api.openai.com/v1/models")
            
            # Check that the URL was rewritten in the prepared request
            call_args = mock_send.call_args
            self.assertIsNotNone(call_args)
            # HTTPAdapter.send receives a PreparedRequest
            prepared_request = call_args[0][0] if call_args[0] else None
            self.assertIsNotNone(prepared_request)
            self.assertEqual(prepared_request.url, "http://localhost:8080/openai/v1/models")
    
    def test_no_rewrite_for_unmatched_domain(self):
        """Test that unmatched domains are not rewritten."""
        rules = {
            "api.openai.com": "http://localhost:8080/openai"
        }
        
        Interceptor.init(routing_rules=rules, verbose=False)
        
        with patch('requests.adapters.HTTPAdapter.send') as mock_send:
            from requests.models import Response
            mock_response = Response()
            mock_response.status_code = 200
            mock_response.reason = 'OK'
            mock_response.headers = {}
            mock_response.url = "https://api.github.com/users"
            mock_response.raw = MagicMock()
            mock_response.raw.release_conn = MagicMock()
            mock_send.return_value = mock_response
            
            session = requests.Session()
            session.get("https://api.github.com/users")
            
            # Check that the URL was NOT rewritten
            call_args = mock_send.call_args
            self.assertIsNotNone(call_args)
            prepared_request = call_args[0][0] if call_args[0] else None
            self.assertIsNotNone(prepared_request)
            # Should remain unchanged
            self.assertEqual(prepared_request.url, "https://api.github.com/users")
    
    def test_context_manager(self):
        """Test the context manager for temporary interception."""
        rules = {
            "api.openai.com": "http://localhost:8080/openai"
        }
        
        with patch('requests.adapters.HTTPAdapter.send') as mock_send:
            from requests.models import Response
            mock_response = Response()
            mock_response.status_code = 200
            mock_response.reason = 'OK'
            mock_response.headers = {}
            mock_response.url = "http://localhost:8080/openai/v1/models"
            mock_response.raw = MagicMock()
            mock_response.raw.release_conn = MagicMock()
            mock_send.return_value = mock_response
            
            with Interceptor.scope(routing_rules=rules, verbose=False):
                session = requests.Session()
                session.get("https://api.openai.com/v1/models")
                
                # Check that the URL was rewritten
                call_args = mock_send.call_args
                self.assertIsNotNone(call_args)
                prepared_request = call_args[0][0] if call_args[0] else None
                self.assertIsNotNone(prepared_request)
                self.assertEqual(prepared_request.url, "http://localhost:8080/openai/v1/models")
            
            # After context, interceptor should be stopped
            self.assertFalse(Interceptor.is_active())
    
    def test_agent_token_injection(self):
        """Test that agent token is added to intercepted requests."""
        import os
        
        # Set agent token
        os.environ["KINDRED_AGENT_TOKEN"] = "test-token-123"
        
        rules = {
            "api.stripe.com": "http://localhost:8080/stripe-mock"
        }
        
        Interceptor.init(routing_rules=rules, verbose=False)
        
        with patch('requests.adapters.HTTPAdapter.send') as mock_send:
            from requests.models import Response
            mock_response = Response()
            mock_response.status_code = 200
            mock_response.reason = 'OK'
            mock_response.headers = {}
            mock_response.url = "http://localhost:8080/stripe-mock/v1/charges"
            mock_response.raw = MagicMock()
            mock_response.raw.release_conn = MagicMock()
            mock_send.return_value = mock_response
            
            session = requests.Session()
            session.get("https://api.stripe.com/v1/charges")
            
            # Check that the request was made with the token header
            call_args = mock_send.call_args
            self.assertIsNotNone(call_args)
            prepared_request = call_args[0][0] if call_args[0] else None
            self.assertIsNotNone(prepared_request)
            
            # Check that the token header was added
            self.assertIn('X-Kindred-Agent-Token', prepared_request.headers)
            self.assertEqual(prepared_request.headers['X-Kindred-Agent-Token'], 'test-token-123')
        
        Interceptor.stop()
        # Clean up
        if "KINDRED_AGENT_TOKEN" in os.environ:
            del os.environ["KINDRED_AGENT_TOKEN"]


if __name__ == '__main__':
    unittest.main()

