graphql-engine/server/tests-py/test_websocket_init_cookie.py

import json
from urllib.parse import urlparse

import websocket
import pytest
from context import PytestConf

if not PytestConf.config.getoption("--test-ws-init-cookie"):
    pytest.skip("--test-ws-init-cookie flag is missing, skipping tests", allow_module_level=True)

def url(hge_ctx):
    ws_url = urlparse(hge_ctx.hge_url)._replace(scheme='ws', path='/v1alpha1/graphql')
    return ws_url.geturl()

@pytest.mark.usefixtures('auth_hook')
class TestWebsocketInitCookie():
    """
    test if cookie is sent when initing the websocket connection, is our auth
    webhook receiving the cookie
    """
    dir = 'queries/remote_schemas'

    @pytest.fixture(autouse=True)
    def transact(self, hge_ctx):
        hge_ctx.v1q_f(self.dir + '/person_table.yaml')
        yield
        hge_ctx.v1q_f(self.dir + '/drop_person_table.yaml')

    def _send_query(self, hge_ctx):
        ws_url = url(hge_ctx)
        headers = {'Cookie': 'foo=bar;'}
        ws = websocket.create_connection(ws_url, header=headers)
        init_payload = {
            'type': 'connection_init',
            'payload': {'headers': {}}
        }
        ws.send(json.dumps(init_payload))
        payload = {
            'type': 'start',
            'id': '1',
            'payload': {'query': 'query { person {name}}'}
        }
        ws.send(json.dumps(payload))
        return ws

    def test_websocket_init_cookie_used(self, hge_ctx):
        if hge_ctx.ws_read_cookie == 'noread':
            pytest.skip('cookie is not to be read')
        ws = self._send_query(hge_ctx)
        it = 0
        while True:
            raw = ws.recv()
            frame = json.loads(raw)
            if frame['type'] == 'data':
                assert 'person' in frame['payload']['data']
                break
            elif it == 10:
                print('max try over')
                assert False
                break
            elif frame['type'] == 'connection_error' or frame['type'] == 'error':
                print(frame)
                assert False
                break
            it = it + 1

    def test_websocket_init_cookie_not_used(self, hge_ctx):
        if hge_ctx.ws_read_cookie == 'read':
            pytest.skip('cookie is read')

        ws = self._send_query(hge_ctx)
        it = 0
        while True:
            raw = ws.recv()
            frame = json.loads(raw)
            if frame['type'] ==  'data':
                print('got data')
                assert False
                break
            elif it == 10:
                print('max try over')
                assert False
                break
            elif frame['type'] == 'connection_error':
                print(frame)
                assert frame['payload'] == 'Authentication hook unauthorized this request'
                break
            elif frame['type'] == 'error':
                print(frame)
                assert False
                break
            it = it + 1