From 3ba94e02da1db516aecb9d1ab77cb76c74b8e127 Mon Sep 17 00:00:00 2001
From: inexcode <inex.code@selfprivacy.org>
Date: Mon, 24 Oct 2022 03:55:48 +0300
Subject: [PATCH] Test of the redis subscription

---
 selfprivacy_api/graphql/schema.py | 22 +++++++++++++++++++---
 shell.nix                         |  2 ++
 2 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/selfprivacy_api/graphql/schema.py b/selfprivacy_api/graphql/schema.py
index dff9304..8e46981 100644
--- a/selfprivacy_api/graphql/schema.py
+++ b/selfprivacy_api/graphql/schema.py
@@ -2,6 +2,9 @@
 # pylint: disable=too-few-public-methods
 
 import asyncio
+import async_timeout
+import redis.asyncio as redis
+
 from typing import AsyncGenerator
 import strawberry
 from selfprivacy_api.graphql import IsAuthenticated
@@ -90,9 +93,22 @@ class Subscription:
 
     @strawberry.subscription(permission_classes=[IsAuthenticated])
     async def count(self, target: int = 100) -> AsyncGenerator[int, None]:
-        for i in range(target):
-            yield i
-            await asyncio.sleep(0.5)
+        r = redis.from_url('unix:///run/redis-sp-api/redis.sock')
+        pubsub = r.pubsub()
+        await pubsub.psubscribe("__keyspace@0__:api_test")
+        while True:
+            try:
+                async with async_timeout.timeout(1):
+                    message = await pubsub.get_message()
+                    if message:
+                        if message['data'] == 'set':
+                            await r.get('api_test')
+                            yield int(await r.get('api_test'))
+                    else:
+                        await asyncio.sleep(0.01)
+            except asyncio.TimeoutError:
+                pass
+
 
 
 schema = strawberry.Schema(query=Query, mutation=Mutation, subscription=Subscription)
diff --git a/shell.nix b/shell.nix
index 0ccb99d..219ff18 100644
--- a/shell.nix
+++ b/shell.nix
@@ -16,6 +16,7 @@ let
     typing-extensions
     psutil
     black
+    redis
     fastapi
     uvicorn
     (buildPythonPackage rec {
@@ -33,6 +34,7 @@ let
         pydantic
         pygments
         poetry
+        redis
         # flask-cors
         (buildPythonPackage rec {
           pname = "graphql-core";