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)