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)