From 8b525ee05be22891a1c18d5e96153807f79930be Mon Sep 17 00:00:00 2001
From: timotheyca <tim@ongoteam.yaconnect.com>
Date: Wed, 22 Dec 2021 19:58:41 +0300
Subject: [PATCH] request_token query_mode+tail

---
 v6d1tokens/client.py       | 17 +++++++++++++----
 v6d1tokens/test-request.py |  8 ++------
 2 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/v6d1tokens/client.py b/v6d1tokens/client.py
index 57a7bb5..cc6b186 100644
--- a/v6d1tokens/client.py
+++ b/v6d1tokens/client.py
@@ -3,21 +3,30 @@ from typing import Optional
 
 import aiohttp
 from v6d0auth import certs
-from v6d0auth.client import mycert
+from v6d0auth.client import mycert, with_role
 
 from v6d1tokens.config import taurl
 
 __all__ = ('request_token',)
 
 
-async def request_token(token_id: str, role: Optional[str] = None) -> str:
+async def request_token(query: str, query_mode: Optional[str] = None, tail: str = '') -> str:
     async with aiohttp.ClientSession() as session:
-        if role is None:
+        if query_mode is None:
+            if tail:
+                raise ValueError('tail is used only for non-None query modes')
             request = session.post(
                 f'{taurl}/get',
-                data=json.dumps([token_id, (await mycert()).hex()]).encode()
+                data=json.dumps([query, (await mycert()).hex()]).encode()
             )
         else:
+            if query_mode == 'token':
+                token_id, role = query, f'token:{query}::{tail}'
+            elif query_mode == 'role':
+                token_id, role = f'role:{query}::{tail}', query
+            else:
+                raise ValueError('unknown query mode')
+            await with_role(role)
             request = session.post(
                 f'{taurl}/get',
                 data=json.dumps([token_id, certs.vkey.encode().hex()]).encode(),
diff --git a/v6d1tokens/test-request.py b/v6d1tokens/test-request.py
index 1a25d39..d944742 100644
--- a/v6d1tokens/test-request.py
+++ b/v6d1tokens/test-request.py
@@ -1,16 +1,12 @@
 import asyncio
 
-from v6d0auth.client import with_role
-
 from v6d1tokens.client import request_token
 
 
 async def main():
     print(await request_token('test'))
-    await with_role('t3st')
-    await with_role('token:te5t::tes7')
-    print(await request_token('role:t3st::7est', 't3st'))
-    print(await request_token('te5t', 'token:te5t::tes7'))
+    print(await request_token('t3st', 'role', '7est'))
+    print(await request_token('te5t', 'token', 'tes7'))
 
 
 if __name__ == '__main__':