mcp/tests/test_redis_token_store.py

import json
import os
import tempfile
import unittest

from app import GatewayApp
from config import GatewayConfig
from services.auth_client import AuthClient
from services.token_store import RedisTokenStore


class FakeRedisClient:
    def __init__(self):
        self.values = {}
        self.expiries = {}
        self.deleted = []

    def set(self, key, value, ex=None):
        self.values[key] = value
        self.expiries[key] = ex
        return True

    def get(self, key):
        return self.values.get(key)

    def delete(self, key):
        self.deleted.append(key)
        self.values.pop(key, None)
        return 1


class DummyTransport:
    def __init__(self):
        self.calls = []

    def post_json(self, url, payload, headers, timeout):
        self.calls.append(
            {
                'url': url,
                'payload': payload,
                'headers': headers,
                'timeout': timeout,
            }
        )
        return {
            'code': 'MCP_0000',
            'msg': 'success',
            'data': {
                'mcp_token': 'MT_session',
                'expire_time': '2099-01-01T00:00:00',
            },
        }


class RedisTokenStoreTest(unittest.TestCase):
    def test_gateway_config_reads_redis_store_settings_from_dotenv(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            dotenv_path = os.path.join(tmp_dir, '.env')
            with open(dotenv_path, 'w', encoding='utf-8') as file:
                file.write('FMS_AUTH_BASE=http://auth.example.test\n')
                file.write('FMS_TOOLS_BASE=http://tools.example.test\n')
                file.write('FMS_TOKEN_STORE=redis\n')
                file.write('FMS_REDIS_HOST=192.168.1.241\n')
                file.write('FMS_REDIS_PORT=6380\n')
                file.write('FMS_REDIS_DB=2\n')
                file.write('FMS_REDIS_PASSWORD=secret\n')
                file.write('FMS_REDIS_PREFIX=fms:mcp:workbuddy:\n')
                file.write('FMS_SESSION_KEY=pc-a:user-a\n')

            config = GatewayConfig.from_env(env={}, dotenv_path=dotenv_path)

            self.assertEqual('redis', config.token_store_type)
            self.assertEqual('192.168.1.241', config.redis_host)
            self.assertEqual(6380, config.redis_port)
            self.assertEqual(2, config.redis_db)
            self.assertEqual('secret', config.redis_password)
            self.assertEqual('fms:mcp:workbuddy:', config.redis_prefix)
            self.assertEqual('pc-a:user-a', config.session_key)

    def test_redis_token_store_isolates_tokens_by_session_key(self):
        redis = FakeRedisClient()
        first = RedisTokenStore(redis, prefix='fms:mcp:workbuddy:', session_key='pc-a:user-a')
        second = RedisTokenStore(redis, prefix='fms:mcp:workbuddy:', session_key='pc-b:user-b')

        first.save('MT_first', '2099-01-01T00:00:00')
        second.save('MT_second', '2099-01-01T00:00:00')

        self.assertEqual('MT_first', first.get()['token'])
        self.assertEqual('MT_second', second.get()['token'])
        self.assertIn('fms:mcp:workbuddy:pc-a:user-a', redis.values)
        self.assertIn('fms:mcp:workbuddy:pc-b:user-b', redis.values)
        self.assertNotEqual(redis.values['fms:mcp:workbuddy:pc-a:user-a'], redis.values['fms:mcp:workbuddy:pc-b:user-b'])

        first.clear()

        self.assertIsNone(first.get())
        self.assertEqual('MT_second', second.get()['token'])
        self.assertEqual(['fms:mcp:workbuddy:pc-a:user-a'], redis.deleted)

    def test_auth_client_refresh_persists_token_and_sends_bearer(self):
        transport = DummyTransport()
        redis = FakeRedisClient()
        store = RedisTokenStore(redis, prefix='fms:mcp:workbuddy:', session_key='pc-a:user-a')
        client = AuthClient(
            base_url='http://auth.example.test',
            client_type='workbuddy',
            token_store=store,
            transport=transport,
            timeout=9,
            session_key='pc-a:user-a',
        )

        client.refresh('MT_old')

        self.assertEqual('http://auth.example.test/mcp/auth/refresh', transport.calls[0]['url'])
        self.assertEqual({'mcp_token': 'MT_old'}, transport.calls[0]['payload'])
        self.assertEqual({'Authorization': 'Bearer MT_old'}, transport.calls[0]['headers'])
        self.assertEqual('MT_session', store.get()['token'])

    def test_gateway_from_config_uses_redis_token_store_when_configured(self):
        config = GatewayConfig(
            auth_base_url='http://auth.example.test',
            tools_base_url='http://tools.example.test',
            token_store_type='redis',
            redis_host='127.0.0.1',
            redis_port=6379,
            redis_db=0,
            redis_password='',
            redis_prefix='fms:mcp:workbuddy:',
            session_key='pc-a:user-a',
        )

        app = GatewayApp.from_config(config, redis_client=FakeRedisClient())

        self.assertIsInstance(app.token_store, RedisTokenStore)
        self.assertEqual('fms:mcp:workbuddy:pc-a:user-a', app.token_store.key)
        self.assertEqual('pc-a:user-a', app.auth_client.session_key)


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