package com.rabbitmq.client.test;
import com.rabbitmq.client.AMQP;
import com.rabbitmq.client.impl.Frame;
import com.rabbitmq.client.impl.nio.ByteBufferInputStream;
import com.rabbitmq.client.impl.nio.ByteBufferOutputStream;
import org.junit.Test;
import java.io.ByteArrayOutputStream;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.ReadableByteChannel;
import java.nio.channels.WritableByteChannel;
import java.util.*;
import static org.hamcrest.Matchers.equalTo;
import static org.junit.Assert.assertThat;
/**
*
*/
public class FrameTest {
@Test public void readFrames() throws IOException {
Random random = new Random();
int nbOfFrames = 100;
AccumulatorReadableByteChannel channel = new AccumulatorReadableByteChannel();
for(int i = 0; i < nbOfFrames; i++) {
byte[] payload = new byte[random.nextInt(2000) + 1];
Frame frame = new Frame(AMQP.FRAME_METHOD, 1, payload);
channel.add(frame);
}
ByteBuffer buffer = ByteBuffer.allocate(8192);
DataInputStream inputStream = new DataInputStream(
new ByteBufferInputStream(channel, buffer)
);
int nbReadFrames = 0;
channel.read(buffer);
buffer.flip();
while(buffer.hasRemaining()) {
Frame.readFrom(inputStream);
nbReadFrames++;
if(!buffer.hasRemaining()) {
buffer.clear();
channel.read(buffer);
buffer.flip();
}
}
assertThat(nbReadFrames, equalTo(nbOfFrames));
}
@Test public void readLargeFrame() throws IOException {
AccumulatorReadableByteChannel channel = new AccumulatorReadableByteChannel();
int [] framesSize = new int [] {100, 75, 20000, 150};
for (int frameSize : framesSize) {
Frame frame = new Frame(AMQP.FRAME_METHOD, 1, new byte[frameSize]);
channel.add(frame);
}
ByteBuffer buffer = ByteBuffer.allocate(8192);
DataInputStream inputStream = new DataInputStream(
new ByteBufferInputStream(channel, buffer)
);
int nbReadFrames = 0;
channel.read(buffer);
buffer.flip();
while(buffer.hasRemaining()) {
Frame.readFrom(inputStream);
nbReadFrames++;
if(!buffer.hasRemaining()) {
buffer.clear();
channel.read(buffer);
buffer.flip();
}
}
assertThat(nbReadFrames, equalTo(framesSize.length));
}
@Test
public void writeFrames() throws IOException {
List<Frame> frames = new ArrayList<Frame>();
Random random = new Random();
int totalFrameSize = 0;
for(int i = 0; i < 100; i++) {
byte[] payload = new byte[random.nextInt(2000) + 1];
Frame frame = new Frame(AMQP.FRAME_METHOD, 1, payload);
frames.add(frame);
totalFrameSize += frame.size();
}
AccumulatorWritableByteChannel channel = new AccumulatorWritableByteChannel();
ByteBuffer buffer = ByteBuffer.allocate(8192);
for (Frame frame : frames) {
frame.writeTo(new DataOutputStream(new ByteBufferOutputStream(channel, buffer)));
}
drain(channel, buffer);
checkWrittenChunks(totalFrameSize, channel);
}
@Test public void writeLargeFrame() throws IOException {
List<Frame> frames = new ArrayList<Frame>();
int totalFrameSize = 0;
int [] framesSize = new int [] {100, 75, 20000, 150};
for (int frameSize : framesSize) {
Frame frame = new Frame(AMQP.FRAME_METHOD, 1, new byte[frameSize]);
frames.add(frame);
totalFrameSize += frame.size();
}
AccumulatorWritableByteChannel channel = new AccumulatorWritableByteChannel();
ByteBuffer buffer = ByteBuffer.allocate(8192);
for (Frame frame : frames) {
frame.writeTo(new DataOutputStream(new ByteBufferOutputStream(channel, buffer)));
}
drain(channel, buffer);
checkWrittenChunks(totalFrameSize, channel);
}
private void checkWrittenChunks(int totalFrameSize, AccumulatorWritableByteChannel channel) {
int totalWritten = 0;
for (byte[] chunk : channel.chunks) {
totalWritten += chunk.length;
}
assertThat(totalWritten, equalTo(totalFrameSize));
}
private static class AccumulatorWritableByteChannel implements WritableByteChannel {
List<byte[]> chunks = new ArrayList<byte[]>();
Random random = new Random();
@Override
public int write(ByteBuffer src) throws IOException {
int remaining = src.remaining();
if(remaining > 0) {
int toRead = random.nextInt(remaining) + 1;
byte [] chunk = new byte[toRead];
src.get(chunk);
chunks.add(chunk);
return toRead;
} else {
return remaining;
}
}
@Override
public boolean isOpen() {
return false;
}
@Override
public void close() throws IOException {
}
}
private static class AccumulatorReadableByteChannel implements ReadableByteChannel {
private List<Byte> bytesOfFrames = new LinkedList<Byte>();
@Override
public int read(ByteBuffer dst) throws IOException {
int remaining = dst.remaining();
int read = 0;
if(remaining > 0) {
Iterator<Byte> iterator = bytesOfFrames.iterator();
while(iterator.hasNext() && read < remaining) {
dst.put(iterator.next());
iterator.remove();
read++;
}
}
return read;
}
@Override
public boolean isOpen() {
return false;
}
@Override
public void close() throws IOException {
}
void add(Frame frame) throws IOException {
ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream(frame.size());
DataOutputStream outputStream = new DataOutputStream(byteArrayOutputStream);
frame.writeTo(outputStream);
outputStream.flush();
for (byte b : byteArrayOutputStream.toByteArray()) {
bytesOfFrames.add(b);
}
}
}
public static void drain(WritableByteChannel channel, ByteBuffer buffer) throws IOException {
buffer.flip();
while(buffer.hasRemaining() && channel.write(buffer) != -1);
buffer.clear();
}
}