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")