def test_should_limit_msg_size(self):
# Set up a 60MB ForwardMsg string
large_msg = _create_dataframe_msg([1, 2, 3])
large_msg.delta.new_element.markdown.body = "X" * 60 * 1000 * 1000
# Create a copy, since serialize_forward_msg modifies the original proto
large_msg_copy = ForwardMsg()
large_msg_copy.CopyFrom(large_msg)
deserialized_msg = ForwardMsg()
deserialized_msg.ParseFromString(serialize_forward_msg(large_msg_copy))
# The metadata should be the same, but contents should be replaced
self.assertEqual(deserialized_msg.metadata, large_msg.metadata)
self.assertNotEqual(deserialized_msg, large_msg)
expected = "Data of size 60.0MB exceeds write limit of 50.0MB"
self.assertEqual(deserialized_msg.delta.new_element.exception.message, expected)
def test_should_limit_msg_size(self):
max_message_size_mb = 50
# Set max message size to defined value
from streamlit.server import server_util
server_util._max_message_size_bytes = None # Reset cached value
config._set_option("server.maxMessageSize", max_message_size_mb,
"test")
# Set up a larger than limit ForwardMsg string
large_msg = _create_dataframe_msg([1, 2, 3])
large_msg.delta.new_element.markdown.body = (
"X" * (max_message_size_mb + 10) * 1000 * 1000)
# Create a copy, since serialize_forward_msg modifies the original proto
large_msg_copy = ForwardMsg()
large_msg_copy.CopyFrom(large_msg)
deserialized_msg = ForwardMsg()
deserialized_msg.ParseFromString(serialize_forward_msg(large_msg_copy))
# The metadata should be the same, but contents should be replaced
self.assertEqual(deserialized_msg.metadata, large_msg.metadata)
self.assertNotEqual(deserialized_msg, large_msg)
self.assertTrue("exceeds the message size limit" in
deserialized_msg.delta.new_element.exception.message)
def read_forward_msg(self, ws_client):
"""Parse the next message from a Websocket client into a ForwardMsg."""
data = yield ws_client.read_message()
message = ForwardMsg()
message.ParseFromString(data)
raise gen.Return(message)
def _parse_msg(msg_string):
"""Parse a ForwardMsg from a string"""
msg = ForwardMsg()
msg.ParseFromString(msg_string)
return msg