def test_message_callback(self):
"""Test that when a callback is included in a Message(), it is invoked
after the message is sent"""
client = mock.Mock()
client._fqdn = "test_server"
client.device_plugins = mock.Mock()
client.device_plugins.get_plugins = mock.Mock(
return_value=['test_plugin'])
client.boot_time = IMLDateTime.utcnow()
client.start_time = IMLDateTime.utcnow()
callback = mock.Mock()
# Disable poll() so that it's not trying to set up sessions, just doing passthrough of messages
with mock.patch("chroma_agent.agent_client.HttpWriter.poll"):
try:
writer = HttpWriter(client)
writer.start()
message = Message("DATA",
"test_plugin", {'key1': 'val1'},
'session_foo',
666,
callback=callback)
writer.put(message)
TIMEOUT = 2
i = 0
while True:
if client.post.call_count and callback.call_count:
break
else:
time.sleep(1)
i += 1
if i > TIMEOUT:
raise RuntimeError(
"Timeout waiting for .post() and callback (%s %s)"
%
(client.post.call_count, callback.call_count))
# Should have sent back the result
self.assertEqual(client.post.call_count, 1)
self.assertDictEqual(
client.post.call_args[0][0], {
'messages': [message.dump(client._fqdn)],
'server_boot_time': client.boot_time.isoformat() + "Z",
'client_start_time':
client.start_time.isoformat() + "Z"
})
# Should have invoked the callback
self.assertEqual(callback.call_count, 1)
finally:
writer.stop()
writer.join()
def inject_messages(*args, **kwargs):
# A control plane message
writer.put(
Message("SESSION_CREATE_REQUEST", "plugin_fuz", None, None,
None))
low_body = DevicePluginMessage('low', PRIO_LOW)
normal_body = DevicePluginMessage('normal', PRIO_NORMAL)
high_body = DevicePluginMessage('high', PRIO_HIGH)
writer.put(Message("DATA", "plugin_foo", low_body, "foo", 0))
writer.put(Message("DATA", "plugin_bar", normal_body, "foo", 1))
writer.put(Message("DATA", "plugin_baz", high_body, "foo", 2))
def test_oversized_messages(self):
"""
Test that oversized messages are dropped and the session is terminated
"""
# Monkey-patch this setting to a lower limit to make testing easier
MAX_BYTES_PER_POST = 1024
from chroma_agent.agent_client import MAX_BYTES_PER_POST as LARGE_MAX_BYTES_PER_POST
def set_post_limit(size):
import chroma_agent.agent_client
chroma_agent.agent_client.MAX_BYTES_PER_POST = size
self.addCleanup(set_post_limit, LARGE_MAX_BYTES_PER_POST)
set_post_limit(MAX_BYTES_PER_POST)
client = mock.Mock()
client._fqdn = "test_server"
client.boot_time = IMLDateTime.utcnow()
client.start_time = IMLDateTime.utcnow()
writer = HttpWriter(client)
def fake_post(envelope):
if len(json.dumps(envelope)) > MAX_BYTES_PER_POST:
daemon_log.info("fake_post(): rejecting oversized message")
raise HttpError()
client.post = mock.Mock(side_effect=fake_post)
TestPlugin = mock.Mock()
mock_plugin_instance = mock.Mock()
mock_plugin_instance.start_session = mock.Mock(
return_value={'foo': 'bar'})
client.device_plugins.get = mock.Mock(
return_value=lambda (plugin_name): mock_plugin_instance)
client.device_plugins.get_plugins = mock.Mock(
return_value={'test_plugin': TestPlugin})
client.sessions = SessionTable(client)
daemon_log.setLevel(logging.DEBUG)
import string
from random import choice
oversized_string = "".join(
choice(string.printable) for i in range(MAX_BYTES_PER_POST))
# There should be one message to set up the session
writer.poll('test_plugin')
self.assertTrue(writer.send())
self.assertEqual(client.post.call_count, 1)
messages = client.post.call_args[0][0]['messages']
self.assertEqual(len(messages), 1)
self.assertEqual(messages[0]['type'], "SESSION_CREATE_REQUEST")
# Pretend we got a SESSION_CREATE_RESPONSE
client.sessions.create('test_plugin', 'id_foo')
self.assertEqual(len(client.sessions._sessions), 1)
# Inject a normal and an oversized message
normal_body = DevicePluginMessage('normal', PRIO_NORMAL)
oversized_body = DevicePluginMessage(oversized_string, PRIO_NORMAL)
writer.put(Message("DATA", "test_plugin", normal_body, "id_foo", 0))
writer.put(Message("DATA", "test_plugin", oversized_body, "id_foo", 1))
# Only the normal message should get through
self.assertTrue(writer.send())
self.assertEqual(client.post.call_count, 2)
messages = client.post.call_args[0][0]['messages']
self.assertEqual(len(messages), 1)
self.assertEqual(messages[0]['type'], "DATA")
# The oversized message should be dropped and the session
# terminated
self.assertFalse(writer.send())
self.assertEqual(client.post.call_count, 3)
self.assertEqual(len(client.sessions._sessions), 0)
# However, we should eventually get a new session for the
# offending plugin
writer.poll('test_plugin')
self.assertTrue(writer.send())
self.assertEqual(client.post.call_count, 4)
messages = client.post.call_args[0][0]['messages']
self.assertEqual(len(messages), 1)
self.assertEqual(messages[0]['type'], "SESSION_CREATE_REQUEST")