def create_subscription_context(instance):
ws = mock.Mock()
yaml_paths = [file_relative_path(__file__, "./workspace.yaml")]
with WorkspaceProcessContext(
instance=instance,
workspace_load_target=WorkspaceFileTarget(paths=yaml_paths),
version="",
read_only=True,
) as process_context:
yield GeventConnectionContext(ws, process_context)
def handle(self, ws, request_context=None):
connection_context = GeventConnectionContext(ws, request_context)
self.on_open(connection_context)
while True:
try:
if connection_context.closed:
raise ConnectionClosedException()
message = connection_context.receive()
except ConnectionClosedException:
self.on_close(connection_context)
return
if message:
msg = json.loads(message)
if msg.get('type', '') == 'start':
Audit.create_audit_entry(msg['payload']['query'],
operation_type='subscription')
self.on_message(connection_context, message)
def test_receive(self):
ws = mock.Mock()
connection_context = GeventConnectionContext(ws=ws)
connection_context.receive()
ws.receive.assert_called()
def test_close(self):
ws = mock.Mock()
connection_context = GeventConnectionContext(ws=ws)
connection_context.close(123)
ws.close.assert_called_with(123)
def test_send_closed(self):
ws = mock.Mock()
ws.closed = True
connection_context = GeventConnectionContext(ws=ws)
connection_context.send("test")
assert not ws.send.called
def test_send(self):
ws = mock.Mock()
ws.closed = False
connection_context = GeventConnectionContext(ws=ws)
connection_context.send("test")
ws.send.assert_called_with("test")