package zmq;

import java.util.Arrays;
import static org.hamcrest.CoreMatchers.is;
import static org.hamcrest.CoreMatchers.not;
import static org.hamcrest.CoreMatchers.notNullValue;
import static org.junit.Assert.assertThat;
import org.junit.Test;

/**
 * This test does one setup, and runs a few tests on that setup.
 *
 * @author lbo
 */
public class TestReqCorrelateRelaxed
{
    private static final String PAYLOAD = "Payload";

    /**
     * Prepares sockets and runs actual tests.
     *
     * Doing it this way so order is guaranteed.
     *
     * @throws Exception
     */
    @Test
    public void overallSetup() throws Exception
    {
        Ctx ctx = ZMQ.init(1);
        assertThat(ctx, notNullValue());

        SocketBase dealer = ZMQ.socket(ctx, ZMQ.ZMQ_DEALER);
        assertThat(dealer, notNullValue());
        boolean brc = ZMQ.bind(dealer, "inproc://a");
        assertThat(brc, is(true));

        SocketBase reqClient = ZMQ.socket(ctx, ZMQ.ZMQ_REQ);
        assertThat(reqClient, notNullValue());

        reqClient.setSocketOpt(ZMQ.ZMQ_REQ_CORRELATE, 1);
        reqClient.setSocketOpt(ZMQ.ZMQ_REQ_RELAXED, 1);

        brc = ZMQ.connect(reqClient, "inproc://a");
        assertThat(brc, is(true));

        byte[] origRequestId = testReqSentFrames(dealer, reqClient);
        testReqRecvGoodRequestId(dealer, reqClient, origRequestId);
    }

    /**
     * Tests that the correct frames are sent by the REQ socket.
     *
     * @param dealer
     * @param reqClient
     * @throws Exception
     * @returns the request ID that was received
     */
    public byte[] testReqSentFrames(SocketBase dealer, SocketBase reqClient) throws Exception
    {
        // Send simple payload over REQ socket
        Msg request = new Msg(PAYLOAD.getBytes());

        assertThat(ZMQ.sendMsg(reqClient, request, 0), is(PAYLOAD.getBytes().length));

        // Verify that the right frames were sent: [request ID, empty frame, payload]
        // 1. request ID
        Msg receivedReqId = ZMQ.recv(dealer, 0);

        assertThat(receivedReqId, notNullValue());
        assertThat(receivedReqId.size(), is(Req.REQUEST_ID_LENGTH));
        assertThat(receivedReqId.flags() & ZMQ.ZMQ_MORE, not(0));

        byte[] buf = new byte[128];
        int requestIdLen = receivedReqId.getBytes(0, buf, 0, 128);
        assertThat(requestIdLen, is(Req.REQUEST_ID_LENGTH));

        byte[] requestId = Arrays.copyOf(buf, Req.REQUEST_ID_LENGTH);

        // 2. empty frame
        Msg receivedEmpty = ZMQ.recv(dealer, 0);

        assertThat(receivedEmpty, notNullValue());
        assertThat(receivedEmpty.size(), is(0));
        assertThat(receivedEmpty.flags() & ZMQ.ZMQ_MORE, not(0));

        // 3. Payload
        Msg receivedPayload = ZMQ.recv(dealer, 0);

        assertThat(receivedPayload, notNullValue());
        assertThat(receivedPayload.size(), is(PAYLOAD.getBytes().length));
        assertThat(receivedPayload.flags() & ZMQ.ZMQ_MORE, is(0));

        int receivedPayloadLen = receivedPayload.getBytes(0, buf, 0, 128);
        assertThat(receivedPayloadLen, is(PAYLOAD.getBytes().length));

        for (int i = 0; i < receivedPayloadLen; i++) {
            assertThat(PAYLOAD.getBytes()[i], is(PAYLOAD.getBytes()[i]));
        }

        return requestId;
    }

    /**
     * Test that REQ sockets with CORRELATE/RELAXED receive a response with
     * correct request ID correctly.
     *
     * @param dealer
     * @param reqClient
     * @param origRequestId the request ID that was sent by the REQ socket
     * earlier
     * @throws Exception
     */
    public void testReqRecvGoodRequestId(SocketBase dealer, SocketBase reqClient, byte[] origRequestId) throws Exception
    {
        Msg requestId = new Msg(origRequestId);
        Msg empty = new Msg();
        Msg responsePayload = new Msg(PAYLOAD.getBytes());

        // Send response
        assertThat(ZMQ.send(dealer, requestId, ZMQ.ZMQ_SNDMORE), is(6));
        assertThat(ZMQ.send(dealer, empty, ZMQ.ZMQ_SNDMORE), is(0));
        assertThat(ZMQ.send(dealer, responsePayload, 0), is(responsePayload.size()));

        // Receive response (payload only)
        Msg receivedResponsePayload = reqClient.recv(0);
        assertThat(receivedResponsePayload, notNullValue());

        byte[] buf = new byte[128];
        int payloadLen = receivedResponsePayload.getBytes(0, buf, 0, 128);
        assertThat(payloadLen, is(PAYLOAD.getBytes().length));
    }
}

Add tests for Req/Relaxed sockets