test_auth_client.py 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161
  1. import os
  2. import tempfile
  3. import unittest
  4. from config import GatewayConfig
  5. from services.auth_client import AuthClient
  6. from services.token_store import InMemoryTokenStore
  7. class DummyTransport:
  8. def __init__(self):
  9. self.calls = []
  10. self.responses = []
  11. def post_json(self, url, payload, headers, timeout):
  12. self.calls.append(
  13. {
  14. 'url': url,
  15. 'payload': payload,
  16. 'headers': headers,
  17. 'timeout': timeout,
  18. }
  19. )
  20. if self.responses:
  21. return self.responses.pop(0)
  22. if url.endswith('/mcp/auth/refresh'):
  23. return {
  24. 'code': 'MCP_0000',
  25. 'msg': 'success',
  26. 'data': {
  27. 'mcp_token': 'MT_refresh',
  28. 'expire_time': '2099-01-02T00:00:00',
  29. },
  30. }
  31. return {
  32. 'code': 'MCP_0000',
  33. 'msg': 'success',
  34. 'data': {},
  35. }
  36. class AuthClientTest(unittest.TestCase):
  37. def test_gateway_config_reads_expected_env_keys(self):
  38. env = {
  39. 'MCP_AUTH_BASE_URL': 'http://base.example.test',
  40. 'MCP_TOOLS_BASE_URL': 'http://tools.example.test',
  41. 'MCP_CLIENT_TYPE': 'workbuddy',
  42. 'MCP_TIMEOUT_SECONDS': '9',
  43. 'MCP_REFRESH_SKEW_SECONDS': '120',
  44. }
  45. config = GatewayConfig.from_env(env, dotenv_path=os.path.join(tempfile.gettempdir(), 'missing-auth-client.env'))
  46. self.assertEqual('http://base.example.test', config.auth_base_url)
  47. self.assertEqual('http://tools.example.test', config.tools_base_url)
  48. self.assertEqual('workbuddy', config.client_type)
  49. self.assertEqual(9, config.timeout_seconds)
  50. self.assertEqual(120, config.refresh_skew_seconds)
  51. def test_auth_client_uses_token_routes_and_updates_token_store(self):
  52. transport = DummyTransport()
  53. store = InMemoryTokenStore(refresh_skew_seconds=60)
  54. client = AuthClient(
  55. base_url='http://base.example.test',
  56. client_type='workbuddy',
  57. token_store=store,
  58. transport=transport,
  59. timeout=9,
  60. )
  61. refresh = client.refresh('MT_exchange')
  62. self.assertEqual('MT_refresh', store.get()['token'])
  63. revoke = client.revoke('MT_refresh')
  64. self.assertEqual('MT_refresh', refresh['data']['mcp_token'])
  65. self.assertIsNone(store.get())
  66. self.assertEqual({}, revoke['data'])
  67. self.assertEqual('http://base.example.test/mcp/auth/refresh', transport.calls[0]['url'])
  68. self.assertEqual('http://base.example.test/mcp/auth/revoke', transport.calls[1]['url'])
  69. self.assertEqual({'mcp_token': 'MT_exchange'}, transport.calls[0]['payload'])
  70. self.assertEqual({'mcp_token': 'MT_refresh'}, transport.calls[1]['payload'])
  71. def test_refresh_and_revoke_send_bearer_header_and_revoke_clears_token(self):
  72. transport = DummyTransport()
  73. store = InMemoryTokenStore(refresh_skew_seconds=60)
  74. store.save('MT_exchange', '2099-01-01T00:00:00')
  75. client = AuthClient(
  76. base_url='http://base.example.test',
  77. client_type='workbuddy',
  78. token_store=store,
  79. transport=transport,
  80. timeout=9,
  81. )
  82. client.refresh('MT_exchange')
  83. client.revoke('MT_refresh')
  84. self.assertEqual('Bearer MT_exchange', transport.calls[0]['headers']['Authorization'])
  85. self.assertEqual('Bearer MT_refresh', transport.calls[1]['headers']['Authorization'])
  86. self.assertIsNone(store.get())
  87. def test_revoke_does_not_clear_token_on_failure(self):
  88. transport = DummyTransport()
  89. transport.responses = [
  90. {'code': 'MCP_9999', 'msg': 'revoke failed'}
  91. ]
  92. store = InMemoryTokenStore(refresh_skew_seconds=60)
  93. store.save('MT_token', '2099-01-01T00:00:00')
  94. client = AuthClient(
  95. base_url='http://base.example.test',
  96. client_type='workbuddy',
  97. token_store=store,
  98. transport=transport,
  99. timeout=9,
  100. )
  101. result = client.revoke('MT_token')
  102. self.assertIsNotNone(store.get())
  103. self.assertEqual(result['code'], 'MCP_9999')
  104. def test_refresh_does_not_persist_token_without_expire_time(self):
  105. transport = DummyTransport()
  106. transport.responses = [
  107. {
  108. 'code': 'MCP_0000',
  109. 'data': {
  110. 'mcp_token': 'MT_no_expire'
  111. }
  112. }
  113. ]
  114. store = InMemoryTokenStore(refresh_skew_seconds=60)
  115. client = AuthClient(
  116. base_url='http://base.example.test',
  117. client_type='workbuddy',
  118. token_store=store,
  119. transport=transport,
  120. timeout=9,
  121. session_key='test_session'
  122. )
  123. client.refresh('MT_old')
  124. self.assertIsNone(store.get())
  125. def test_auth_client_no_longer_exposes_exchange(self):
  126. client = AuthClient(
  127. base_url='http://base.example.test',
  128. client_type='workbuddy',
  129. token_store=InMemoryTokenStore(),
  130. )
  131. self.assertFalse(hasattr(client, 'exchange'))
  132. if __name__ == '__main__':
  133. unittest.main()