Repository: Zapix/mtpylon (language: python)

Issue:
Set session id to current context
set session id to current context


Discussion / hints:

Fix patch:
diff --git a/mtpylon/contextvars.py b/mtpylon/contextvars.py
index 5d31a2e..c5a980a 100644
--- a/mtpylon/contextvars.py
+++ b/mtpylon/contextvars.py
@@ -41,6 +41,10 @@
 server_salt_var: ContextVar[long] = ContextVar('server_salt_var')
 
 
+"""Stores session id for current context"""
+session_id_var: ContextVar[long] = ContextVar('session_id_var')
+
+
 """
 Stores message that has been received
 """
diff --git a/mtpylon/message_handler/strategies/handle_rpc_query_message.py b/mtpylon/message_handler/strategies/handle_rpc_query_message.py
index a2e945a..f766c6d 100644
--- a/mtpylon/message_handler/strategies/handle_rpc_query_message.py
+++ b/mtpylon/message_handler/strategies/handle_rpc_query_message.py
@@ -10,6 +10,7 @@
 from mtpylon.contextvars import (
     income_message_var,
     server_salt_var,
+    session_id_var,
 )
 from mtpylon.serialization import CallableFunc
 from mtpylon.middlewares import apply_middleware, MiddleWareFunc, Handler
@@ -80,11 +81,12 @@ async def run_rpc_query(
     result = await handler(request, **value.params)
 
     server_salt = server_salt_var.get()
+    session_id = session_id_var.get()
 
     await sender.send_encrypted_message(
         request,
         server_salt,
-        message.session_id,
+        session_id,
         result,
         True
     )
diff --git a/mtpylon/middlewares/set_session_id.py b/mtpylon/middlewares/set_session_id.py
index c1305a7..d2db268 100644
--- a/mtpylon/middlewares/set_session_id.py
+++ b/mtpylon/middlewares/set_session_id.py
@@ -3,7 +3,7 @@
 
 from aiohttp.web import Request
 
-from mtpylon.contextvars import income_message_var
+from mtpylon.contextvars import income_message_var, session_id_var
 from mtpylon.messages import Message
 
 from .types import Handler
@@ -21,4 +21,6 @@ async def set_session_id(
         if not await session_subject.has_session(income_message.session_id):
             await session_subject.create_session(income_message.session_id)
 
+        session_id_var.set(income_message.session_id)
+
     return await handler(request, **params)


Test patch:
diff --git a/tests/test_message_handler/test_strategies/test_handle_rpc_query_message.py b/tests/test_message_handler/test_strategies/test_handle_rpc_query_message.py
index 046b83c..f86ef52 100644
--- a/tests/test_message_handler/test_strategies/test_handle_rpc_query_message.py
+++ b/tests/test_message_handler/test_strategies/test_handle_rpc_query_message.py
@@ -11,7 +11,7 @@
     run_rpc_query
 )
 from mtpylon.service_schema.constructors import RpcResult, RpcError
-from mtpylon.contextvars import server_salt_var
+from mtpylon.contextvars import server_salt_var, session_id_var
 
 from tests.simpleschema import get_task, set_task, Task
 
@@ -64,6 +64,7 @@ async def test_run_rpc_query_success():
     )
 
     server_salt_var.set(server_salt)
+    session_id_var.set(session_id)
 
     await run_rpc_query([], sender, request, message)
 
@@ -102,6 +103,7 @@ async def test_run_rpc_query_error():
     )
 
     server_salt_var.set(server_salt)
+    session_id_var.set(session_id)
 
     await run_rpc_query([], sender, request, message)
 
@@ -136,6 +138,7 @@ async def test_run_rpc_unexpected_error():
     )
 
     server_salt_var.set(server_salt)
+    session_id_var.set(session_id)
 
     await run_rpc_query([], sender, request, message)
 
diff --git a/tests/test_middlewares/test_session_id.py b/tests/test_middlewares/test_session_id.py
index 9ed5e16..ebe6c4c 100644
--- a/tests/test_middlewares/test_session_id.py
+++ b/tests/test_middlewares/test_session_id.py
@@ -4,7 +4,10 @@
 import pytest
 
 from mtpylon.types import long
-from mtpylon.contextvars import income_message_var
+from mtpylon.contextvars import (
+    income_message_var,
+    session_id_var,
+)
 from mtpylon.serialization import CallableFunc
 from mtpylon.messages import Message
 from mtpylon.sessions import (
@@ -61,6 +64,8 @@ async def test_create_session_id():
     )
     assert awaited_arg == expected_event
 
+    assert session_id_var.get() == session_id
+
 
 @pytest.mark.asyncio
 async def test_session_id_created_before():
@@ -100,3 +105,5 @@ async def test_session_id_created_before():
 
     handler.assert_awaited()
     session_observer.update.assert_not_awaited()
+
+    assert session_id_var.get() == session_id