/** 
 * 
 * Copyright 2004 Michael Gaffney
 * 
 * Licensed under the Apache License, Version 2.0 (the "License"); 
 * you may not use this file except in compliance with the License. 
 * You may obtain a copy of the License at 
 * 
 * http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, 
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
 * See the License for the specific language governing permissions and 
 * limitations under the License. 
 * 
 **/
package org.codehaus.activemq.service.impl;

import org.codehaus.activemq.broker.Broker;
import org.codehaus.activemq.message.ActiveMQXid;
import org.codehaus.activemq.service.TransactionTask;
import org.codehaus.activemq.store.PreparedTransactionStore;
import org.jmock.cglib.Mock;
import org.jmock.cglib.MockObjectTestCase;

import javax.jms.JMSException;
import javax.transaction.xa.XAException;
import javax.transaction.xa.XAResource;
import java.util.Map;

/**
 * @author <a href="mailto:michael.gaffney@panacya.com">Michael Gaffney </a>
 */
public class XATransactionCommandTest extends MockObjectTestCase {

    private XATransactionCommand txnCmd;
    private ActiveMQXid stubXid;
    private Mock mockBroker;
    private Mock mockXaTxnMap;
    private Mock mockTxnStore;
    private Mock mockTxnTask;
    private Mock stubTxnTask;

    public XATransactionCommandTest(String name) {
        setName(name);
    }

    public void testRollbackFromPreparedStateWithPostRollbackFailure() {
        putTxnCmdInPreparedState();
        addPostRollbackFailure();
        expectRemoveXidCalled();
        expectTxnStorageRemoveCalled();
        doRollbackExpectException(XAException.XAER_RMERR);
        endInFinishedState();
    }
        
    public void testRollbackFromPreparedState() {
        putTxnCmdInPreparedState();
        expectPostRollbackCalled();
        expectRemoveXidCalled();
        expectTxnStorageRemoveCalled();
        doRollbackExpectSuccess();
        endInFinishedState();
    }
        
    public void testRollbackFromInUseStateWithPostRollbackFailure() {
        putTxnCmdInUseState();
        addPostRollbackFailure();
        expectRemoveXidCalled();
        doRollbackExpectException(XAException.XAER_RMERR);
        endInFinishedState();
    }
        
    public void testRollbackFromInUseState() {
        putTxnCmdInUseState();
        expectPostRollbackCalled();
        expectRemoveXidCalled();
        doRollbackExpectSuccess();
        endInFinishedState();
    }
        
    public void testRollbackFromStartState() {
        putTxnCmdInStartState();
        expectRemoveXidCalled();
        doRollbackExpectSuccess();
        endInFinishedState();
    }
        
    public void testPrepareFromPrepareState() {
        putTxnCmdInPreparedState();
        expectRemoveXidNeverCalled();
        doPrepareExpectException(XAException.XAER_PROTO);
        endInPreparedState();
    }

    public void testPrepareFromInUseStateWithPrePrepareFailure() {
        putTxnCmdInUseState();
        addPrePrepareFailure();        
        expectPostRollbackCalled();
        expectRemoveXidCalled();
        doPrepareExpectException(XAException.XA_RBOTHER);
        endInFinishedState();
    }

    public void testPrepareFromInUseState() {
        putTxnCmdInUseState();
        expectPrePrepareCalled();        
        expectTxnStoragePutCalled();
        doPrepareExpectSuccess(XAResource.XA_OK);
        endInPreparedState();
    }

    public void testPrepareFromStartState() {
        putTxnCmdInStartState();
        expectRemoveXidCalled();
        doPrepareExpectSuccess(XAResource.XA_RDONLY);
        endInFinishedState();
    }
    
    public void testCommitTwoPhaseFromPreparedStateWithPostCommitFailure() {
        putTxnCmdInPreparedState();
        doPostCommitFailure(false);
    }

    public void testCommitTwoPhaseFromPreparedState() {
        putTxnCmdInPreparedState();
        expectPostCommitCalled();
        doCommitExpectSuccess(false);
        endInFinishedState();
    }

    public void testCommitOnePhaseFromInUseStateWithPostCommitFailure() {
        putTxnCmdInUseState();
        expectPrePrepareCalled();
        doPostCommitFailure(true);
    }

    public void testCommitOnePhaseFromInUseStateWithPrePrepareFailure() {
        putTxnCmdInUseState();
        addPrePrepareFailure();
        expectPostRollbackCalled();
        expectRemoveXidCalled();
        doCommitExpectException(true, XAException.XA_RBOTHER);
        endInFinishedState();
    }

    public void testCommitOnePhaseFromInUseState() {
        putTxnCmdInUseState();
        expectPrePrepareCalled();
        expectPostCommitCalled();
        doCommitExpectSuccess(true);
        endInFinishedState();
    }

    public void testCommitTwoPhaseFromInUseState() {
        putTxnCmdInUseState();
        expectRemoveXidNeverCalled();
        doCommitExpectException(false, XAException.XAER_PROTO);
        endInUseState();
    }

    public void testCommitOnePhaseFromStartState() {
        putTxnCmdInStartState();
        doCommitExpectSuccess(true);
        endInFinishedState();
    }

    public void testCommitTwoPhaseFromStartState() {
        putTxnCmdInStartState();
        expectRemoveXidNeverCalled();
        doCommitExpectException(false, XAException.XAER_PROTO);
        endInStartState();
    }

    public void testCommitTwoPhaseFromFinishedState() {
        putTxnCmdInFinishedState();
        expectRemoveXidNeverCalled();
        doCommitExpectException(false, XAException.XAER_PROTO);
        endInFinishedState();
    }

    public void testPrePrepare() throws Throwable {
        putTxnCmdInStartState();
        expectPrePrepareCalled();
        assertTxnCmdInUseState();   
        txnCmd.prePrepare();
        endInUseState();
    }

    public void testAddPrePrepare() throws Throwable {
        putTxnCmdInStartState();
        txnCmd.addPrePrepareTask((TransactionTask) mockTxnTask.proxy());
        endInUseState();
    }

    public void testAddPostCommitTask() throws Throwable {
        putTxnCmdInStartState();
        txnCmd.addPostCommitTask((TransactionTask) mockTxnTask.proxy());
        endInUseState();
    }

    public void testAddPostRollbackTask() throws Throwable {
        putTxnCmdInStartState();
        txnCmd.addPostRollbackTask((TransactionTask) mockTxnTask.proxy());
        endInUseState();
    }

    protected void setUp() throws Exception {
        setupStubs();
        setupMocks();
        setupTxnCommand();
    }

    private void setupStubs() throws JMSException {
        stubXid = new ActiveMQXid("1:2:3");
        stubTxnTask = new Mock(TransactionTask.class);        
    }

    private void setupTxnCommand() {
        txnCmd = new XATransactionCommand((Broker) mockBroker.proxy(),
                stubXid, (Map) mockXaTxnMap.proxy(),
                (PreparedTransactionStore) mockTxnStore.proxy());
    }

    private void setupMocks() {
        mockBroker = new Mock(Broker.class);
        mockXaTxnMap = new Mock(Map.class);
        mockTxnStore = new Mock(PreparedTransactionStore.class);
        mockTxnTask = new Mock(TransactionTask.class);
    }

    private void verifyMocks() {
        mockBroker.verify();
        mockXaTxnMap.verify();
        mockTxnStore.verify();
        mockTxnTask.verify();
    }

    private void doRollbackExpectSuccess() {
        try {
            txnCmd.rollback();
        } catch (XAException e) {
            fail("No XAException should be thrown");
        }
    }
        
    private void doRollbackExpectException(final int expectedErrorCode) {
        try {
            txnCmd.rollback();
            fail("An XAException should have been thrown");
        } catch (XAException e) {
            assertEquals("Incorrect XAException error code", expectedErrorCode, e.errorCode);
        }
    }
    
    private void doPrepareExpectException(final int expectedErrorCode) {
        try {
            txnCmd.prepare();
            fail("An XAException should have been thrown");
        } catch (XAException e) {
            assertEquals("Incorrect XAException error code", expectedErrorCode, e.errorCode);
        }
    }
    
    private void doPrepareExpectSuccess(final int expectedXaCode) {
        try {
            assertEquals("Incorrect XAResource code returned.", expectedXaCode, txnCmd.prepare());
        } catch (XAException e) {
            fail("No XAException should be thrown");
        }
    }
    
    private void doPostCommitFailure(final boolean onePhase) {
        addPostCommitFailure();
        expectRemoveXidCalled();
        doCommitExpectException(onePhase, XAException.XAER_RMERR);
        endInFinishedState();
    }

    private void doCommitExpectSuccess(final boolean onePhase) {
        expectRemoveXidCalled();
        try {
            txnCmd.commit(onePhase);
        } catch (XAException e) {
            fail("No XAException should be thrown");
        }
    }

    private void doCommitExpectException(final boolean onePhase, final int expectedErrorCode) {
        try {
            txnCmd.commit(onePhase);
            fail("An XAException should have been thrown");
        } catch (XAException e) {
            assertEquals("Incorrect XAException error code", expectedErrorCode, e.errorCode);
        }
    }

    private void addPostRollbackFailure() {
        stubTxnTask.expects(once()).method("execute").will(throwException(new Throwable("Stub")));
        txnCmd.addPostRollbackTask((TransactionTask) stubTxnTask.proxy());
    }

    private void addPostCommitFailure() {
        stubTxnTask.expects(once()).method("execute").will(throwException(new Throwable("Stub")));
        txnCmd.addPostCommitTask((TransactionTask) stubTxnTask.proxy());
    }

    private void addPrePrepareFailure() {
        stubTxnTask.expects(once()).method("execute").will(throwException(new Throwable("Stub")));
        txnCmd.addPrePrepareTask((TransactionTask) stubTxnTask.proxy());
    }

    private void expectPrePrepareCalled() {
        mockTxnTask.expects(once()).method("execute").with(same(mockBroker.proxy()));
        txnCmd.addPrePrepareTask((TransactionTask) mockTxnTask.proxy());
    }

    private void expectPostCommitCalled() {
        mockTxnTask.expects(once()).method("execute").with(same(mockBroker.proxy()));
        txnCmd.addPostCommitTask((TransactionTask) mockTxnTask.proxy());
    }

    private void expectPostRollbackCalled() {
        mockTxnTask.expects(once()).method("execute").with(same(mockBroker.proxy()));
        txnCmd.addPostRollbackTask((TransactionTask) mockTxnTask.proxy());
    }

    private void expectRemoveXidCalled() {
        mockXaTxnMap.expects(once()).method("remove").with(same(stubXid));
    }

    private void expectRemoveXidNeverCalled() {
        mockXaTxnMap.expects(never()).method("remove");
    }
    
    private void expectTxnStoragePutCalled() {
        mockTxnStore.expects(once()).method("put").with(same(stubXid), same(txnCmd));
    }
    
    private void expectTxnStorageRemoveCalled() {
        mockTxnStore.expects(once()).method("remove").with(same(stubXid));
    }
    
    private void putTxnCmdInStartState() {
        txnCmd.setState(AbstractTransaction.START_STATE);
        assertTxnCmdInStartState();
    }

    private void putTxnCmdInUseState() {
        txnCmd.setState(AbstractTransaction.IN_USE_STATE);
        assertTxnCmdInUseState();
    }

    private void putTxnCmdInPreparedState() {
        txnCmd.setState(AbstractTransaction.PREPARED_STATE);
        assertTxnCmdInPreparedState();
    }

    private void putTxnCmdInFinishedState() {
        txnCmd.setState(AbstractTransaction.FINISHED_STATE);
        assertTxnCmdInFinishedState();
    }
    
    private void assertTxnCmdInStartState() {
        assertEquals("Not in START_STATE", AbstractTransaction.START_STATE, txnCmd.getState());
    }

    private void assertTxnCmdInFinishedState() {
        assertEquals("Not in FINISHED_STATE", AbstractTransaction.FINISHED_STATE, txnCmd.getState());
    }

    private void assertTxnCmdInUseState() {
        assertEquals("Not in IN_USE_STATE", AbstractTransaction.IN_USE_STATE, txnCmd.getState());
    }

    private void assertTxnCmdInPreparedState() {
        assertEquals("Not in PREPARED_STATE", AbstractTransaction.PREPARED_STATE, txnCmd.getState());
    }

    private void endInFinishedState() {
        assertTxnCmdInFinishedState();
        verifyMocks();
    }
    
    private void endInUseState() {
        assertTxnCmdInUseState();
        verifyMocks();
    }

    private void endInStartState() {
        assertTxnCmdInStartState();
        verifyMocks();
    }

    private void endInPreparedState() {
        assertTxnCmdInPreparedState();
        verifyMocks();
    }
}
