def testSubmitRequestsWithContextManager(self, messageBusConnectorClassMock):
    requests = [
      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 = "foofar"

    messageBusConnectorMock = messageBusConnectorClassMock.return_value

    with ModelSwapperInterface() as interface:
      batchID = interface.submitRequests(modelID=modelID, requests=requests)

      modelMQName = interface._modelInputQueueNamePrefix + modelID
      notificationMQName = interface._schedulerNotificationQueueName

    # Verify
    self.assertIsInstance(batchID, str)

    msg = RequestMessagePackager.marshal(
      batchID=batchID,
      batchState=BatchPackager.marshal(batch=requests))

    self.assertEqual(messageBusConnectorMock.publish.call_count, 2)

    messageBusConnectorMock.publish.assert_any_call(
      modelMQName, msg, persistent=True)

    messageBusConnectorMock.publish.assert_any_call(
      notificationMQName, json.dumps(modelID), persistent=False)
  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 testMarshalAndUnmarshal(self):
    requestBatch = [
        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"]),
      ]
    batchState = BatchPackager.marshal(batch=requestBatch)
    msg = RequestMessagePackager.marshal(batchID="foobar",
                                         batchState=batchState)

    r = RequestMessagePackager.unmarshal(msg)

    self.assertEqual(r.batchID, "foobar")
    self.assertEqual(r.batchState, batchState)

    # Make sure we aren't forgetting to test any returned fields
    self.assertEqual(set(["batchID", "batchState"]), set(r._fields))