def testConsumeModelSchedulerNotifications(self,
messageBusConnectorClassMock):
modelID = "foobar"
msg = json.dumps(modelID)
messageBusConnectorMock = messageBusConnectorClassMock.return_value
ackMock = Mock(return_value=None)
messageBusConnectorMock.consume.return_value = Mock(
spec_set=message_bus_connector._QueueConsumer,
__iter__=lambda *args, **kwargs: iter(
[message_bus_connector._ConsumedMessage(body=msg, ack=ackMock)]))
with ModelSwapperInterface() as interface:
self.assertEqual(len(interface._consumers), 0)
consumer = interface.consumeModelSchedulerNotifications()
self.assertIsNotNone(consumer._mqConsumer)
self.assertEqual(len(interface._consumers), 1)
notification = next(iter(consumer))
self.assertEqual(notification.value, modelID)
self.assertTrue(callable(notification.ack))
self.assertIs(notification.ack, ackMock)
consumer.close()
self.assertIsNone(consumer._mqConsumer)
self.assertEqual(len(interface._consumers), 0)
self.assertEqual(len(interface._consumers), 0)
def testConsumeRequestsNonBlocking(self, messageBusConnectorClassMock):
expectedRequests = (
ModelCommand(commandID="abc", method="defineModel",
args={'key1': 4098, 'key2': 4139}),
ModelInputRow(rowID="foo", data=[1, 2, "Sep 21 02:24:21 UTC 2013"]),
ModelInputRow(rowID="bar", data=[9, 54, "Sep 21 02:24:38 UTC 2013"]),
)
modelID = "foobar"
batchID = uuid.uuid1().hex
msg = RequestMessagePackager.marshal(
batchID=batchID,
batchState=BatchPackager.marshal(batch=expectedRequests))
messageBusConnectorMock = messageBusConnectorClassMock.return_value
ackMock = Mock(return_value=None)
messageBusConnectorMock.consume.return_value = Mock(
spec_set=message_bus_connector._QueueConsumer,
__iter__=lambda *args, **kwargs: iter(
[message_bus_connector._ConsumedMessage(body=msg, ack=ackMock)]))
with ModelSwapperInterface() as interface:
self.assertEqual(len(interface._consumers), 0)
with interface.consumeRequests(modelID, blocking=False) as consumer:
self.assertIsNotNone(consumer._mqConsumer)
self.assertEqual(len(interface._consumers), 1)
batch = next(iter(consumer))
self.assertEqual(batch.batchID, batchID)
self.assertEqual(batch.objects, expectedRequests)
self.assertTrue(callable(batch.ack))
self.assertIs(batch.ack, ackMock)
# Make sure we didn't skip any fields
self.assertEqual(set(batch._fields), set(["batchID", "objects", "ack"]))
self.assertIsNone(consumer._mqConsumer)
self.assertEqual(len(interface._consumers), 0)
self.assertEqual(len(interface._consumers), 0)
modelMQName = interface._modelInputQueueNamePrefix + modelID
messageBusConnectorMock.consume.assert_called_once_with(
modelMQName, blocking=False)
def testContextManagerAndConsumeResults(self, messageBusConnectorClassMock):
expectedResults = (
ModelCommandResult(commandID="abc", method="testMethod", status=0,
args={'key1': 4098, 'key2': 4139}),
ModelInferenceResult(rowID="foo", status=0, anomalyScore=1.3),
ModelInferenceResult(rowID="bar", status=0, anomalyScore=2.9)
)
modelID = "foobar"
msg = ResultMessagePackager.marshal(
modelID=modelID,
batchState=BatchPackager.marshal(batch=expectedResults))
messageBusConnectorMock = messageBusConnectorClassMock.return_value
ackMock = Mock(return_value=None)
messageBusConnectorMock.consume.return_value = Mock(
spec_set=message_bus_connector._QueueConsumer,
__iter__=lambda *args, **kwargs: iter(
[message_bus_connector._ConsumedMessage(
body=msg, ack=ackMock)]))
with ModelSwapperInterface() as interface:
self.assertEqual(len(interface._consumers), 0)
consumer = interface.consumeResults()
self.assertIsNotNone(consumer._mqConsumer)
self.assertEqual(len(interface._consumers), 1)
batch = next(iter(consumer))
self.assertEqual(batch.modelID, modelID)
self.assertEqual(batch.objects, expectedResults)
self.assertTrue(callable(batch.ack))
self.assertIs(batch.ack, ackMock)
self.assertEqual(set(batch._fields), set(["modelID", "objects", "ack"]))
consumer.close()
self.assertIsNone(consumer._mqConsumer)
self.assertEqual(len(interface._consumers), 0)
self.assertEqual(len(interface._consumers), 0)