mcp/tests/test_gateway_session_store.py

import unittest

from services.gateway_session_store import GatewaySessionStore


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

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

    def get(self, key):
        item = self.values.get(key)
        return None if item is None else item['value']

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


class GatewaySessionStoreTest(unittest.TestCase):
    def test_save_and_get_session_by_hashed_gateway_session_id(self):
        redis = FakeRedis()
        store = GatewaySessionStore(redis, prefix='fms:mcp:gateway:', ttl_seconds=3600)

        store.save('GWS_employee_a', {
            'mcp_token': 'MT_A',
            'admin_id': 1,
            'company_id': 10,
        })
        session = store.get('GWS_employee_a')

        self.assertEqual('MT_A', session['mcp_token'])
        self.assertEqual(1, session['admin_id'])
        self.assertNotIn('GWS_employee_a', list(redis.values.keys())[0])
        self.assertEqual(3600, list(redis.values.values())[0]['ex'])

    def test_two_gateway_sessions_do_not_overlap(self):
        redis = FakeRedis()
        store = GatewaySessionStore(redis, prefix='fms:mcp:gateway:', ttl_seconds=3600)

        store.save('GWS_employee_a', {'mcp_token': 'MT_A'})
        store.save('GWS_employee_b', {'mcp_token': 'MT_B'})

        self.assertEqual('MT_A', store.get('GWS_employee_a')['mcp_token'])
        self.assertEqual('MT_B', store.get('GWS_employee_b')['mcp_token'])

    def test_delete_removes_only_current_session(self):
        redis = FakeRedis()
        store = GatewaySessionStore(redis, prefix='fms:mcp:gateway:', ttl_seconds=3600)

        store.save('GWS_employee_a', {'mcp_token': 'MT_A'})
        store.save('GWS_employee_b', {'mcp_token': 'MT_B'})
        store.delete('GWS_employee_a')

        self.assertIsNone(store.get('GWS_employee_a'))
        self.assertEqual('MT_B', store.get('GWS_employee_b')['mcp_token'])


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