selfprivacy-rest-api/tests/test_graphql/test_websocket.py

229 lines
7.2 KiB
Python
Raw Normal View History

# from selfprivacy_api.graphql.subscriptions.jobs import JobSubscriptions
import pytest
import asyncio
2024-05-27 20:21:11 +00:00
from typing import Generator
from time import sleep
from starlette.testclient import WebSocketTestSession
2024-05-15 18:15:16 +00:00
from selfprivacy_api.jobs import Jobs
2024-05-27 20:21:11 +00:00
from selfprivacy_api.actions.api_tokens import TOKEN_REPO
from selfprivacy_api.graphql import IsAuthenticated
from tests.conftest import DEVICE_WE_AUTH_TESTS_WITH
from tests.test_jobs import jobs as empty_jobs
2024-05-27 20:21:11 +00:00
# We do not iterate through them yet
TESTED_SUBPROTOCOLS = ["graphql-transport-ws"]
2024-05-15 18:15:16 +00:00
2024-05-15 20:43:17 +00:00
JOBS_SUBSCRIPTION = """
jobUpdates {
uid
typeId
name
description
status
statusText
progress
createdAt
updatedAt
finishedAt
error
result
}
"""
2024-05-15 18:15:16 +00:00
2024-05-15 11:29:20 +00:00
2024-05-27 20:21:11 +00:00
def connect_ws_authenticated(authorized_client) -> WebSocketTestSession:
token = "Bearer " + str(DEVICE_WE_AUTH_TESTS_WITH["token"])
return authorized_client.websocket_connect(
"/graphql",
subprotocols=TESTED_SUBPROTOCOLS,
params={"token": token},
)
def connect_ws_not_authenticated(client) -> WebSocketTestSession:
return client.websocket_connect(
"/graphql",
subprotocols=TESTED_SUBPROTOCOLS,
params={"token": "I like vegan icecream but it is not a valid token"},
)
def init_graphql(websocket):
websocket.send_json({"type": "connection_init", "payload": {}})
ack = websocket.receive_json()
assert ack == {"type": "connection_ack"}
@pytest.fixture
def authenticated_websocket(
authorized_client,
) -> Generator[WebSocketTestSession, None, None]:
# We use authorized_client only tohave token in the repo, this client by itself is not enough to authorize websocket
ValueError(TOKEN_REPO.get_tokens())
with connect_ws_authenticated(authorized_client) as websocket:
yield websocket
@pytest.fixture
def unauthenticated_websocket(client) -> Generator[WebSocketTestSession, None, None]:
with connect_ws_not_authenticated(client) as websocket:
yield websocket
2024-05-15 11:29:20 +00:00
def test_websocket_connection_bare(authorized_client):
2024-05-15 18:15:16 +00:00
client = authorized_client
with client.websocket_connect(
"/graphql", subprotocols=["graphql-transport-ws", "graphql-ws"]
) as websocket:
2024-05-15 11:29:20 +00:00
assert websocket is not None
assert websocket.scope is not None
2024-05-15 18:15:16 +00:00
def test_websocket_graphql_init(authorized_client):
2024-05-15 18:15:16 +00:00
client = authorized_client
with client.websocket_connect(
"/graphql", subprotocols=["graphql-transport-ws"]
) as websocket:
websocket.send_json({"type": "connection_init", "payload": {}})
ack = websocket.receive_json()
assert ack == {"type": "connection_ack"}
def test_websocket_graphql_ping(authorized_client):
client = authorized_client
with client.websocket_connect(
"/graphql", subprotocols=["graphql-transport-ws"]
) as websocket:
2024-05-15 18:31:16 +00:00
# https://github.com/enisdenjo/graphql-ws/blob/master/PROTOCOL.md#ping
websocket.send_json({"type": "ping", "payload": {}})
pong = websocket.receive_json()
assert pong == {"type": "pong"}
2024-05-15 18:15:16 +00:00
def api_subscribe(websocket, id, subscription):
websocket.send_json(
{
"id": id,
"type": "subscribe",
"payload": {
"query": "subscription TestSubscription {" + subscription + "}",
},
}
)
2024-05-15 20:43:17 +00:00
def test_websocket_subscription_minimal(authorized_client):
# Test a small endpoint that exists specifically for tests
2024-05-15 20:43:17 +00:00
client = authorized_client
with client.websocket_connect(
"/graphql", subprotocols=["graphql-transport-ws"]
) as websocket:
init_graphql(websocket)
arbitrary_id = "3aaa2445"
api_subscribe(websocket, arbitrary_id, "count")
2024-05-15 20:43:17 +00:00
response = websocket.receive_json()
assert response == {
"id": arbitrary_id,
2024-05-15 20:43:17 +00:00
"payload": {"data": {"count": 0}},
"type": "next",
}
response = websocket.receive_json()
assert response == {
"id": arbitrary_id,
2024-05-15 20:43:17 +00:00
"payload": {"data": {"count": 1}},
"type": "next",
}
response = websocket.receive_json()
assert response == {
"id": arbitrary_id,
2024-05-15 20:43:17 +00:00
"payload": {"data": {"count": 2}},
"type": "next",
}
def test_websocket_subscription_minimal_unauthorized(unauthenticated_websocket):
websocket = unauthenticated_websocket
init_graphql(websocket)
arbitrary_id = "3aaa2445"
api_subscribe(websocket, arbitrary_id, "count")
response = websocket.receive_json()
assert response == {
"id": arbitrary_id,
"payload": [{"message": IsAuthenticated.message}],
"type": "error",
}
async def read_one_job(websocket):
# Bug? We only get them starting from the second job update
# That's why we receive two jobs in the list them
# The first update gets lost somewhere
response = websocket.receive_json()
return response
@pytest.mark.asyncio
2024-05-27 20:21:11 +00:00
async def test_websocket_subscription(authenticated_websocket, event_loop, empty_jobs):
websocket = authenticated_websocket
init_graphql(websocket)
arbitrary_id = "3aaa2445"
api_subscribe(websocket, arbitrary_id, JOBS_SUBSCRIPTION)
2024-05-27 20:21:11 +00:00
future = asyncio.create_task(read_one_job(websocket))
jobs = []
jobs.append(Jobs.add("bogus", "bogus.bogus", "yyyaaaaayy it works"))
sleep(0.5)
jobs.append(Jobs.add("bogus2", "bogus.bogus", "yyyaaaaayy it works"))
response = await future
data = response["payload"]["data"]
jobs_received = data["jobUpdates"]
received_names = [job["name"] for job in jobs_received]
for job in jobs:
assert job.name in received_names
assert len(jobs_received) == 2
for job in jobs:
for api_job in jobs_received:
if (job.name) == api_job["name"]:
assert api_job["uid"] == str(job.uid)
assert api_job["typeId"] == job.type_id
assert api_job["name"] == job.name
assert api_job["description"] == job.description
assert api_job["status"] == job.status
assert api_job["statusText"] == job.status_text
assert api_job["progress"] == job.progress
assert api_job["createdAt"] == job.created_at.isoformat()
assert api_job["updatedAt"] == job.updated_at.isoformat()
assert api_job["finishedAt"] == None
assert api_job["error"] == None
assert api_job["result"] == None
def test_websocket_subscription_unauthorized(unauthenticated_websocket):
websocket = unauthenticated_websocket
init_graphql(websocket)
id = "3aaa2445"
api_subscribe(websocket, id, JOBS_SUBSCRIPTION)
2024-05-27 20:21:11 +00:00
response = websocket.receive_json()
# I do not really know why strawberry gives more info on this
# One versus the counter
payload = response["payload"][0]
assert isinstance(payload, dict)
assert "locations" in payload.keys()
# It looks like this 'locations': [{'column': 32, 'line': 1}]
# We cannot test locations feasibly
del payload["locations"]
2024-05-27 20:21:11 +00:00
assert response == {
"id": id,
"payload": [{"message": IsAuthenticated.message, "path": ["jobUpdates"]}],
2024-05-27 20:21:11 +00:00
"type": "error",
}