mcp/tests/test_auth_client.py

import os
import os
import tempfile
import unittest

from config import GatewayConfig
from services.auth_client import AuthClient
from services.token_store import InMemoryTokenStore


class DummyTransport:
    def __init__(self):
        self.calls = []
        self.responses = []  # Queue of responses to return

    def post_json(self, url, payload, headers, timeout):
        self.calls.append(
            {
                'url': url,
                'payload': payload,
                'headers': headers,
                'timeout': timeout,
            }
        )

        # If custom responses are queued, return the next one
        if self.responses:
            return self.responses.pop(0)

        # Default responses
        if url.endswith('/mcp/auth/exchange'):
            return {
                'code': 'MCP_0000',
                'msg': 'success',
                'data': {
                    'mcp_token': 'MT_exchange',
                    'expire_time': '2099-01-01T00:00:00',
                },
            }
        if url.endswith('/mcp/auth/refresh'):
            return {
                'code': 'MCP_0000',
                'msg': 'success',
                'data': {
                    'mcp_token': 'MT_refresh',
                    'expire_time': '2099-01-02T00:00:00',
                },
            }
        return {
            'code': 'MCP_0000',
            'msg': 'success',
            'data': {},
        }


class AuthClientTest(unittest.TestCase):
    def test_gateway_config_reads_expected_env_keys(self):
        env = {
            'MCP_AUTH_BASE_URL': 'http://base.example.test',
            'MCP_TOOLS_BASE_URL': 'http://tools.example.test',
            'MCP_CLIENT_TYPE': 'workbuddy',
            'MCP_TIMEOUT_SECONDS': '9',
            'MCP_REFRESH_SKEW_SECONDS': '120',
        }

        config = GatewayConfig.from_env(env, dotenv_path=os.path.join(tempfile.gettempdir(), 'missing-auth-client.env'))

        self.assertEqual('http://base.example.test', config.auth_base_url)
        self.assertEqual('http://tools.example.test', config.tools_base_url)
        self.assertEqual('workbuddy', config.client_type)
        self.assertEqual(9, config.timeout_seconds)
        self.assertEqual(120, config.refresh_skew_seconds)

    def test_auth_client_uses_auth_routes_and_updates_token_store(self):
        transport = DummyTransport()
        store = InMemoryTokenStore(refresh_skew_seconds=60)
        client = AuthClient(
            base_url='http://base.example.test',
            client_type='workbuddy',
            token_store=store,
            transport=transport,
            timeout=9,
        )

        exchange = client.exchange('AUTH123')
        refresh = client.refresh('MT_exchange')
        self.assertEqual('MT_refresh', store.get()['token'])
        revoke = client.revoke('MT_refresh')

        self.assertEqual('MT_exchange', exchange['data']['mcp_token'])
        self.assertEqual('MT_refresh', refresh['data']['mcp_token'])
        self.assertIsNone(store.get())
        self.assertEqual({}, revoke['data'])
        self.assertEqual('http://base.example.test/mcp/auth/exchange', transport.calls[0]['url'])
        self.assertEqual('http://base.example.test/mcp/auth/refresh', transport.calls[1]['url'])
        self.assertEqual('http://base.example.test/mcp/auth/revoke', transport.calls[2]['url'])
        self.assertEqual({'auth_code': 'AUTH123', 'client_type': 'workbuddy'}, transport.calls[0]['payload'])
        self.assertEqual({'mcp_token': 'MT_exchange'}, transport.calls[1]['payload'])
        self.assertEqual({'mcp_token': 'MT_refresh'}, transport.calls[2]['payload'])

    def test_refresh_and_revoke_send_bearer_header_and_revoke_clears_token(self):
        transport = DummyTransport()
        store = InMemoryTokenStore(refresh_skew_seconds=60)
        store.save('MT_exchange', '2099-01-01T00:00:00')
        client = AuthClient(
            base_url='http://base.example.test',
            client_type='workbuddy',
            token_store=store,
            transport=transport,
            timeout=9,
        )

        client.refresh('MT_exchange')
        client.revoke('MT_refresh')

        self.assertEqual('Bearer MT_exchange', transport.calls[0]['headers']['Authorization'])
        self.assertEqual('Bearer MT_refresh', transport.calls[1]['headers']['Authorization'])
        self.assertIsNone(store.get())

    def test_revoke_does_not_clear_token_on_failure(self):
        """Test that revoke does not clear token when response code is not MCP_0000"""
        transport = DummyTransport()
        transport.responses = [
            {'code': 'MCP_9999', 'msg': 'revoke failed'}
        ]
        store = InMemoryTokenStore(refresh_skew_seconds=60)
        store.save('MT_token', '2099-01-01T00:00:00')

        client = AuthClient(
            base_url='http://base.example.test',
            client_type='workbuddy',
            token_store=store,
            transport=transport,
            timeout=9,
        )

        result = client.revoke('MT_token')

        # Token should NOT be cleared on failure
        self.assertIsNotNone(store.get())
        self.assertEqual(result['code'], 'MCP_9999')

    def test_persist_token_handles_missing_expire_time(self):
        """Test that token is not persisted when expire_time is missing"""
        transport = DummyTransport()
        transport.responses = [
            {
                'code': 'MCP_0000',
                'data': {
                    'mcp_token': 'MT_no_expire'
                    # Missing expire_time
                }
            }
        ]
        store = InMemoryTokenStore(refresh_skew_seconds=60)

        client = AuthClient(
            base_url='http://base.example.test',
            client_type='workbuddy',
            token_store=store,
            transport=transport,
            timeout=9,
            session_key='test_session'
        )

        client.exchange('AC_test')

        # Token should not be saved without expire_time
        self.assertIsNone(store.get())


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