✨ Support async sessions

authx/_internal/_memory.py

@@ -19,34 +19,34 @@ def __init__(self) -> None:
         """
         self.raw_memory_store = {}
 
-    def has_session_id(self, session_id: str) -> bool:
+    async def has_session_id(self, session_id: str) -> bool:
         return session_id in self.raw_memory_store
 
-    def has_no_session_id(self, session_id: str) -> bool:
+    async def has_no_session_id(self, session_id: str) -> bool:
         return session_id not in self.raw_memory_store
 
-    def create_store(self, session_id: str) -> Dict[str, Any]:
+    async def create_store(self, session_id: str) -> Dict[str, Any]:
         self.raw_memory_store[session_id] = {
             "created_at": int(time.time()),
             "store": {},
         }
-        self.save_store(session_id)
+        await self.save_store(session_id)
         return self.raw_memory_store.get(session_id, {}).get("store", {})
 
-    def get_store(self, session_id: str) -> Optional[Dict[str, Any]]:
+    async def get_store(self, session_id: str) -> Optional[Dict[str, Any]]:
         if self.raw_memory_store.get(session_id):
             return self.raw_memory_store.get(session_id, {}).get("store")
         else:
             return None
 
-    def save_store(self, session_id: str) -> None:
-        self.get_store(session_id)
+    async def save_store(self, session_id: str) -> None:
+        await self.get_store(session_id)
 
-    def gc(self) -> None:
-        if len(self.raw_memory_store) >= 100:  # pragma: no cover
-            self.cleanup_old_sessions()
+    async def gc(self) -> None:
+        if len(self.raw_memory_store) >= 100:
+            await self.cleanup_old_sessions()
 
-    def cleanup_old_sessions(self) -> None:
+    async def cleanup_old_sessions(self) -> None:
         current_time = int(time.time())
         sessions_to_delete = [
             session_id

authx/_internal/_utils.py

@@ -55,6 +55,7 @@ def to_UTC_without_tz(
     event_timestamp: str, format: str = "%Y-%m-%d %H:%M:%S.%f"
 ) -> str:
     dt = datetime.strptime(event_timestamp, format)
+    dt = dt.replace(tzinfo=tz.utc)
     return dt.astimezone(tz.utc).strftime(format)
 
 

tests/internal/test_memory.py

@@ -10,81 +10,90 @@ def memory_io():
     return MemoryIO()
 
 
-def test_create_store(memory_io):
+@pytest.mark.asyncio
+async def test_create_store(memory_io):
     session_id = "123"
-    store = memory_io.create_store(session_id)
+    store = await memory_io.create_store(session_id)
     assert store == {}
     assert memory_io.raw_memory_store[session_id]["store"] == {}
 
 
-def test_get_store_existing(memory_io):
+@pytest.mark.asyncio
+async def test_get_store_existing(memory_io):
     session_id = "123"
     memory_io.raw_memory_store[session_id] = {
         "created_at": int(time()),
         "store": {"key": "value"},
     }
-    store = memory_io.get_store(session_id)
+    store = await memory_io.get_store(session_id)
     assert store == {"key": "value"}
 
 
-def test_get_store_nonexistent(memory_io):
+@pytest.mark.asyncio
+async def test_get_store_nonexistent(memory_io):
     session_id = "123"
-    store = memory_io.get_store(session_id)
+    store = await memory_io.get_store(session_id)
     assert store is None
 
 
-def test_save_store(memory_io):
+@pytest.mark.asyncio
+async def test_save_store(memory_io):
     session_id = "123"
     memory_io.raw_memory_store[session_id] = {
         "created_at": int(time()),
         "store": {"key": "value"},
     }
-    memory_io.save_store(session_id)
-    assert memory_io.get_store(session_id) == {"key": "value"}
+    await memory_io.save_store(session_id)
+    assert await memory_io.get_store(session_id) == {"key": "value"}
 
 
-def test_cleanup_old_sessions(memory_io):
+@pytest.mark.asyncio
+async def test_cleanup_old_sessions(memory_io):
     current_time = int(time())
     memory_io.raw_memory_store = {
         "1": {"created_at": current_time - 3600 * 12 - 1, "store": {}},
         "2": {"created_at": current_time - 3600 * 12, "store": {}},
         "3": {"created_at": current_time - 3600 * 12 + 1, "store": {}},
     }
-    memory_io.cleanup_old_sessions()
+    await memory_io.cleanup_old_sessions()
     expected_output = {
         "2": {"created_at": current_time - 3600 * 12, "store": {}},
         "3": {"created_at": current_time - 3600 * 12 + 1, "store": {}},
     }
     assert memory_io.raw_memory_store == expected_output
 
 
-def test_has_session_id():
+@pytest.mark.asyncio
+async def test_has_session_id():
     store = MemoryIO()
-    store.create_store("test-id")
-    assert store.has_session_id("test-id")
-    assert not store.has_no_session_id("test-id")
+    await store.create_store("test-id")
+    assert await store.has_session_id("test-id")
+    assert not await store.has_no_session_id("test-id")
 
 
-def test_get_store():
+@pytest.mark.asyncio
+async def test_get_store():
     store = MemoryIO()
-    store.create_store("test-id")
-    assert store.get_store("test-id") == {}
-    assert store.get_store("nonexistent-id") is None
+    await store.create_store("test-id")
+    assert await store.get_store("test-id") == {}
+    assert await store.get_store("nonexistent-id") is None
 
 
-def populate_old_sessions(memory_io, count, created_at):
+@pytest.mark.asyncio
+async def populate_old_sessions(memory_io, count, created_at):
     for i in range(count):
         memory_io.raw_memory_store[str(i)] = {
             "created_at": created_at,
             "store": {},
         }
 
 
-def test_gc_cleanup_old_sessions(memory_io):
+@pytest.mark.asyncio
+async def test_gc_cleanup_old_sessions(memory_io):
     # Populate raw_memory_store with 100 sessions older than 12 hours
     current_time = int(time())
     twelve_hours_ago = current_time - 3600 * 12
-    populate_old_sessions(memory_io, 100, twelve_hours_ago)
+    await populate_old_sessions(memory_io, 100, twelve_hours_ago)
 
     # Add one more session within 12 hours
     extra_session_id = "1000"
@@ -94,7 +103,7 @@ def test_gc_cleanup_old_sessions(memory_io):
     }
 
     # Ensure gc triggers cleanup
-    memory_io.gc()
+    await memory_io.gc()
 
     # Ensure old sessions are cleaned up
     assert len(memory_io.raw_memory_store) == 101