
import os
import unittest
import asyncio
import certifi
from unittest.mock import patch, MagicMock
from coze_workload_identity import requests as coze_requests
from coze_workload_identity import httpx as coze_httpx

class TestPPEHeader(unittest.IsolatedAsyncioTestCase):
    def setUp(self):
        # Reset environment variables before each test
        if 'COZE_ENV_LANE' in os.environ:
            del os.environ['COZE_ENV_LANE']
        
        # Reset default session/clients
        coze_requests._default_session = None
        coze_httpx._default_client = None
        coze_httpx._default_async_client = None

        # Mock proxy and CA env vars to avoid errors
        os.environ['COZE_OUTBOUND_AUTH_PROXY'] = 'http://proxy.example.com'
        os.environ['identity_ticket'] = 'dummy_ticket' # Required for HttpsProxy to work
        # Use a real CA bundle to satisfy httpx's SSL context creation
        os.environ['COZE_OUTBOUND_AUTH_PROXY_CA_PATH'] = certifi.where()

    def tearDown(self):
        if 'COZE_ENV_LANE' in os.environ:
            del os.environ['COZE_ENV_LANE']
        if 'COZE_OUTBOUND_AUTH_PROXY' in os.environ:
            del os.environ['COZE_OUTBOUND_AUTH_PROXY']
        if 'identity_ticket' in os.environ:
            del os.environ['identity_ticket']
        if 'COZE_OUTBOUND_AUTH_PROXY_CA_PATH' in os.environ:
            del os.environ['COZE_OUTBOUND_AUTH_PROXY_CA_PATH']

    def test_requests_ppe_header(self):
        os.environ['COZE_ENV_LANE'] = 'ppe_test'
        session = coze_requests.session()
        self.assertIn('x-tt-env', session.headers)
        self.assertEqual(session.headers['x-tt-env'], 'ppe_test')

    def test_requests_no_ppe_header(self):
        os.environ['COZE_ENV_LANE'] = 'prod'
        session = coze_requests.session()
        self.assertNotIn('x-tt-env', session.headers)

    def test_httpx_client_ppe_header(self):
        os.environ['COZE_ENV_LANE'] = 'ppe_test'
        with coze_httpx.Client() as client:
            self.assertIn('x-tt-env', client.headers)
            self.assertEqual(client.headers['x-tt-env'], 'ppe_test')

    def test_httpx_client_no_ppe_header(self):
        os.environ['COZE_ENV_LANE'] = 'prod'
        with coze_httpx.Client() as client:
            self.assertNotIn('x-tt-env', client.headers)

    async def test_httpx_async_client_ppe_header(self):
        os.environ['COZE_ENV_LANE'] = 'ppe_test'
        async with coze_httpx.AsyncClient() as client:
            self.assertIn('x-tt-env', client.headers)
            self.assertEqual(client.headers['x-tt-env'], 'ppe_test')

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