@Test
  public void testBackupFailure() throws Exception {
    IPartitionManager pm = ManagementContainer.getInstance().getPartitionManager();

    m_part.setAccessibleRemotely(true);
    pm.updatePartition(m_part);
    m_part = pm.getPartition(TEST_PID);

    Customer cust =
        TestUtils.createActiveMailboxTestCustomer(Capabilities.ActiveMailboxSupport.ARCHIVE);
    try {
      MockMessageImporter mock = new MockMessageImporter(getClass().getSimpleName());
      PartitionStoreStage stg =
          new PartitionStoreStage(mock) {
            public void finish() {
              super.finish();
              ((StatComponent) m_importer.getStatComponent()).stop();
            }
          };
      MockStageController ctrl = new MockStageController("test");
      stg.setStageController(ctrl);

      assertEquals(0, mock.getNumFailedMessages());
      assertEquals(0, mock.getNumSucceededMessages());
      // Partition dummyPart = new Partition(10, "sdcvsdv");

      File backupSubDir = new File(m_part.getStorageDirectory(), Partition.BACKUP_SUBDIR);

      Pair<Partition, List<IMessageContext>> buffer =
          new Pair<Partition, List<IMessageContext>>(
              m_part, generateMockMessageContext(10000, 10, cust));

      stg.process(buffer);

      // Nothing should have failed
      assertEquals(0, mock.getNumFailedMessages());
      assertEquals(10, ctrl.getOutputCount());

      assertTrue(backupSubDir.exists());

      // Stop backup writers, delete backup directory and then make sure everything works when we
      // start back up
      Utils.finishBackupWriters();

      FileUtils.deleteDirectory(backupSubDir);

      assertFalse(backupSubDir.exists());
      Utils.startBackupWriters();

      ctrl.resetOutputCount();
      buffer =
          new Pair<Partition, List<IMessageContext>>(
              m_part, generateMockMessageContext(100000, 10, cust));

      stg.process(buffer);

      // Nothing should have failed
      assertEquals(0, mock.getNumFailedMessages());
      assertEquals(10, ctrl.getOutputCount());

      assertTrue(backupSubDir.exists());
      stg.finish();
    } finally {
      TestUtils.quietDeleteCustomer(cust);
    }
  }
  @Test
  public void testFailureTracking() throws Exception {
    final IPartitionProxy partitionProxy = createNiceMock(IPartitionProxy.class);
    final IPartitionProxy partitionProxy2 = createNiceMock(IPartitionProxy.class);
    long msgId = 2L;
    IFailureTrackingManager failMgr = ManagementContainer.getInstance().getFailureTrackingManager();
    StorageImporter importer = new MockStorageImporter();
    Customer customer = TestUtils.createTestCustomer();
    try {
      PartitionStoreStage stage = new PartitionStoreStage(importer);
      Pipeline pipeline =
          PipelineBuilder.start(
                  ManagementContainer.getInstance().getConfiguration(),
                  MessageImporter.PIPELINE_CONFIG_SECTION)
              .singleton("store-stage", stage)
              .get();
      File envelopeFile =
          File.createTempFile(
              "PartitionStoreTest", "test", ManagementContainer.getInstance().getNfsRoot());
      IMessageContext msgCtx = new MessageContext(envelopeFile, null);
      msgCtx.setCustomer(customer);
      msgCtx.setIndexMessage(true);
      msgCtx.setInternalId(msgId);
      msgCtx.setPartition(m_part);
      msgCtx.setPartitionID(m_part.getId());
      Collection<IMessageContext> msgs = Collections.singletonList(msgCtx);
      Collection<IMessageContext> emptyResult = Collections.emptyList();

      expect(partitionProxy.storeMessages(msgs)).andReturn(emptyResult).anyTimes();
      replay(partitionProxy);

      PartitionManager pm =
          new PartitionManager(
              ManagementContainer.getInstance().getPool(ManagementContainer.AM_POOL_NAME)) {
            @Override
            public IPartitionProxy getContentProxy(int id) {
              return partitionProxy;
            }
          };
      PartitionManagerTest.preparePartitionManager(pm);
      stage.setPartitionManager(pm);

      FailureRecord failureRecord =
          FailureTrackingHelper.buildFailureRecord(msgCtx, FailureCategory.STORAGE, null, null);
      failureRecord.setStoredOnPartition(false);
      failMgr.insertMessageFailure(failureRecord);
      msgCtx.setInternalFailureId(msgId);
      pipeline.process(
          new Pair<Partition, List<IMessageContext>>(m_part, Collections.singletonList(msgCtx)));
      failureRecord = failMgr.getFailureRecord(msgId);
      assertTrue("Should find stored on partition to be true", failureRecord.isStoredOnPartition());
      Map<String, String> props = msgCtx.readEnvelope();
      assertTrue(
          "envelope should contain failure ID",
          props.containsKey(ParseEnvelopeStage.INTERNAL_MESSAGE_FAILURE_ID));
      assertEquals(
          "message ID in envelope should match permanent ID",
          String.valueOf(msgId),
          props.get(ParseEnvelopeStage.INTERNAL_MESSAGE_FAILURE_ID));

      // confirm that transient ID gets set back to real ID of message
      // remove the existing failure record
      failMgr.removeFailure(msgId);
      // change the ID
      msgCtx.setInternalFailureId(msgId + 1);
      msgCtx.setInternalId(msgId + 1);
      // insert the new failure record
      failureRecord =
          FailureTrackingHelper.buildFailureRecord(msgCtx, FailureCategory.STORAGE, null, null);
      failMgr.insertMessageFailure(failureRecord);
      // old record should be gone
      failureRecord = failMgr.getFailureRecord(msgId);
      assertNull("Should no longer find the failure record with permanent ID", failureRecord);
      // check that the new record can be fetched
      failureRecord = failMgr.getFailureRecord(msgId + 1);
      assertNotNull("Should find transient ID failure record", failureRecord);
      msgCtx.appendToEnvelope(ParseEnvelopeStage.INTERNAL_MESSAGE_FAILURE_ID, msgId + 1);
      props = msgCtx.readEnvelope();
      assertTrue(
          "envelope should contain failure ID",
          props.containsKey(ParseEnvelopeStage.INTERNAL_MESSAGE_FAILURE_ID));
      assertEquals(
          "message ID in envelope should match permanent ID",
          String.valueOf(msgId + 1),
          props.get(ParseEnvelopeStage.INTERNAL_MESSAGE_FAILURE_ID));

      // now assign back the permanent ID and process through the pipeline
      msgCtx.setInternalId(msgId);
      pipeline.process(
          new Pair<Partition, List<IMessageContext>>(m_part, Collections.singletonList(msgCtx)));

      failureRecord = failMgr.getFailureRecord(msgId + 1);
      assertNull("Should not find any record based on transient message ID", failureRecord);
      failureRecord = failMgr.getFailureRecord(msgId);
      assertNotNull("Should find failure record for permanent ID", failureRecord);
      assertTrue("Should find stored on partition to be true", failureRecord.isStoredOnPartition());
      // envelope should now reflect permanent ID
      props = msgCtx.readEnvelope();
      assertTrue(
          "envelope should contain failure ID",
          props.containsKey(ParseEnvelopeStage.INTERNAL_MESSAGE_FAILURE_ID));
      assertEquals(
          "message ID in envelope should match permanent ID",
          String.valueOf(msgId),
          props.get(ParseEnvelopeStage.INTERNAL_MESSAGE_FAILURE_ID));

      // confirm that state does not change on failure
      // force an exception on any call to storeMessages
      pm =
          new PartitionManager(
              ManagementContainer.getInstance().getPool(ManagementContainer.AM_POOL_NAME)) {
            @Override
            public IPartitionProxy getContentProxy(int id) {
              return partitionProxy2;
            }
          };
      PartitionManagerTest.preparePartitionManager(pm);
      stage.setPartitionManager(pm);

      expect(partitionProxy2.storeMessages(msgs))
          .andThrow(new IPartitionProxy.ServiceNotAvailableException(new Exception()))
          .anyTimes();
      replay(partitionProxy2);
      failureRecord.setStoredOnPartition(false);
      failMgr.insertMessageFailure(failureRecord);
      pipeline.process(
          new Pair<Partition, List<IMessageContext>>(m_part, Collections.singletonList(msgCtx)));
      failureRecord = failMgr.getFailureRecord(msgId);
      assertNotNull("Should find the failure record for message", failureRecord);
      assertFalse(
          "Should find stored on partition to be true", failureRecord.isStoredOnPartition());
    } finally {
      importer.stop();
      TestUtils.quietDeleteCustomer(customer);
    }
  }