Examples with MockIdFactory com.github.ambry.store.MockIdFactory used on opensource projects

Example 11 with MockIdFactory

use of com.github.ambry.store.MockIdFactory in project ambry by linkedin.

the class MessageFormatSendTest method doSendWriteSingleMessageTest.

/**
 * Helper method for testing single message sends.
 * @param encryptionKey the encryption key to include in the message while writing it.
 * @param expectedEncryptionKey the key expected when reading the sent message.
 */
private void doSendWriteSingleMessageTest(ByteBuffer encryptionKey, ByteBuffer expectedEncryptionKey) throws Exception {
    String serviceId = "serviceId";
    String ownerId = "owner";
    String contentType = "bin";
    short accountId = 10;
    short containerId = 2;
    byte[] blob = TestUtils.getRandomBytes(10000);
    byte[] userMetadata = TestUtils.getRandomBytes(2000);
    StoreKey storeKey = new MockId("012345678910123456789012");
    BlobProperties properties = new BlobProperties(blob.length, serviceId, ownerId, contentType, false, 100, accountId, containerId, encryptionKey != null, null, null, null);
    MessageFormatInputStream putStream;
    MessageFormatRecord.MessageHeader_Format header;
    if (putFormat.equals(PutMessageFormatInputStream.class.getSimpleName())) {
        header = getHeader(new PutMessageFormatInputStream(storeKey, encryptionKey == null ? null : encryptionKey.duplicate(), properties, ByteBuffer.wrap(userMetadata), new ByteBufferInputStream(ByteBuffer.wrap(blob)), blob.length, BlobType.DataBlob));
        putStream = new PutMessageFormatInputStream(storeKey, encryptionKey, properties, ByteBuffer.wrap(userMetadata), new ByteBufferInputStream(ByteBuffer.wrap(blob)), blob.length, BlobType.DataBlob);
    } else {
        header = getHeader(new PutMessageFormatBlobV1InputStream(storeKey, properties, ByteBuffer.wrap(userMetadata), new ByteBufferInputStream(ByteBuffer.wrap(blob)), blob.length, BlobType.DataBlob));
        putStream = new PutMessageFormatBlobV1InputStream(storeKey, properties, ByteBuffer.wrap(userMetadata), new ByteBufferInputStream(ByteBuffer.wrap(blob)), blob.length, BlobType.DataBlob);
    }
    ByteBuffer buf1 = ByteBuffer.allocate((int) putStream.getSize());
    putStream.read(buf1.array());
    ArrayList<ByteBuffer> listbuf = new ArrayList<ByteBuffer>();
    listbuf.add(buf1);
    ArrayList<StoreKey> storeKeys = new ArrayList<StoreKey>();
    storeKeys.add(storeKey);
    MockMessageReadSet readSet = new MockMessageReadSet(listbuf, storeKeys);
    MetricRegistry registry = new MetricRegistry();
    MessageFormatMetrics metrics = new MessageFormatMetrics(registry);
    // get all
    MessageFormatSend send = new MessageFormatSend(readSet, MessageFormatFlags.All, metrics, new MockIdFactory());
    Assert.assertEquals(send.sizeInBytes(), putStream.getSize());
    Assert.assertEquals(1, send.getMessageMetadataList().size());
    Assert.assertEquals(null, send.getMessageMetadataList().get(0));
    ByteBuffer bufresult = ByteBuffer.allocate((int) putStream.getSize());
    WritableByteChannel channel = Channels.newChannel(new ByteBufferOutputStream(bufresult));
    while (!send.isSendComplete()) {
        send.writeTo(channel);
    }
    Assert.assertArrayEquals(buf1.array(), bufresult.array());
    Assert.assertTrue(readSet.isPrefetchInfoCorrect(0, readSet.sizeInBytes(0)));
    // get blob
    send = new MessageFormatSend(readSet, MessageFormatFlags.Blob, metrics, new MockIdFactory());
    long blobRecordSize = putFormat.equals(PutMessageFormatInputStream.class.getSimpleName()) ? MessageFormatRecord.Blob_Format_V2.getBlobRecordSize(blob.length) : MessageFormatRecord.Blob_Format_V1.getBlobRecordSize(blob.length);
    Assert.assertEquals(send.sizeInBytes(), blobRecordSize);
    bufresult.clear();
    channel = Channels.newChannel(new ByteBufferOutputStream(bufresult));
    while (!send.isSendComplete()) {
        send.writeTo(channel);
    }
    for (int i = 0; i < blob.length; i++) {
        Assert.assertEquals(blob[i], bufresult.array()[i + (int) blobRecordSize - MessageFormatRecord.Crc_Size - blob.length]);
    }
    if (expectedEncryptionKey == null) {
        Assert.assertEquals(null, send.getMessageMetadataList().get(0));
    } else {
        Assert.assertEquals(expectedEncryptionKey, send.getMessageMetadataList().get(0).getEncryptionKey());
    }
    Assert.assertTrue(readSet.isPrefetchInfoCorrect(0, readSet.sizeInBytes(0)));
    // get user metadata
    send = new MessageFormatSend(readSet, MessageFormatFlags.BlobUserMetadata, metrics, new MockIdFactory());
    long userMetadataRecordSize = MessageFormatRecord.UserMetadata_Format_V1.getUserMetadataSize(ByteBuffer.wrap(userMetadata));
    Assert.assertEquals(send.sizeInBytes(), userMetadataRecordSize);
    bufresult.clear();
    channel = Channels.newChannel(new ByteBufferOutputStream(bufresult));
    while (!send.isSendComplete()) {
        send.writeTo(channel);
    }
    bufresult.flip();
    // read off the header.
    for (int i = 0; i < userMetadataRecordSize - MessageFormatRecord.Crc_Size - userMetadata.length; i++) {
        bufresult.get();
    }
    verifyBlobUserMetadata(userMetadata, bufresult);
    if (expectedEncryptionKey == null) {
        Assert.assertEquals(null, send.getMessageMetadataList().get(0));
    } else {
        Assert.assertEquals(expectedEncryptionKey, send.getMessageMetadataList().get(0).getEncryptionKey());
    }
    Assert.assertTrue(readSet.isPrefetchInfoCorrect(header.getUserMetadataRecordRelativeOffset(), header.getUserMetadataRecordSize()));
    // get blob properties
    send = new MessageFormatSend(readSet, MessageFormatFlags.BlobProperties, metrics, new MockIdFactory());
    long blobPropertiesRecordSize = MessageFormatRecord.BlobProperties_Format_V1.getBlobPropertiesRecordSize(properties);
    Assert.assertEquals(send.sizeInBytes(), blobPropertiesRecordSize);
    bufresult.clear();
    channel = Channels.newChannel(new ByteBufferOutputStream(bufresult));
    while (!send.isSendComplete()) {
        send.writeTo(channel);
    }
    bufresult.flip();
    // read off the header.
    for (int i = 0; i < blobPropertiesRecordSize - MessageFormatRecord.Crc_Size - BlobPropertiesSerDe.getBlobPropertiesSerDeSize(properties); i++) {
        bufresult.get();
    }
    verifyBlobProperties(properties, BlobPropertiesSerDe.getBlobPropertiesFromStream(new DataInputStream(new ByteBufferInputStream(bufresult))));
    Assert.assertEquals(null, send.getMessageMetadataList().get(0));
    Assert.assertTrue(readSet.isPrefetchInfoCorrect(header.getBlobPropertiesRecordRelativeOffset(), header.getBlobPropertiesRecordSize()));
    // get blob info
    send = new MessageFormatSend(readSet, MessageFormatFlags.BlobInfo, metrics, new MockIdFactory());
    Assert.assertEquals(send.sizeInBytes(), blobPropertiesRecordSize + userMetadataRecordSize);
    bufresult.clear();
    channel = Channels.newChannel(new ByteBufferOutputStream(bufresult));
    while (!send.isSendComplete()) {
        send.writeTo(channel);
    }
    bufresult.flip();
    for (int i = 0; i < blobPropertiesRecordSize - MessageFormatRecord.Crc_Size - BlobPropertiesSerDe.getBlobPropertiesSerDeSize(properties); i++) {
        bufresult.get();
    }
    verifyBlobProperties(properties, BlobPropertiesSerDe.getBlobPropertiesFromStream(new DataInputStream(new ByteBufferInputStream(bufresult))));
    for (int i = 0; i < userMetadataRecordSize - userMetadata.length; i++) {
        bufresult.get();
    }
    verifyBlobUserMetadata(userMetadata, bufresult);
    if (expectedEncryptionKey == null) {
        Assert.assertEquals(null, send.getMessageMetadataList().get(0));
    } else {
        Assert.assertEquals(expectedEncryptionKey, send.getMessageMetadataList().get(0).getEncryptionKey());
    }
    Assert.assertTrue(readSet.isPrefetchInfoCorrect(header.getBlobPropertiesRecordRelativeOffset(), header.getBlobPropertiesRecordSize() + header.getUserMetadataRecordSize()));
}

Example 12 with MockIdFactory

use of com.github.ambry.store.MockIdFactory in project ambry by linkedin.

the class MessageFormatSendTest method doSendWriteCompositeMessagesTest.

/**
 * Helper method to test multiple messages in a single Send involving different combinations of header format
 * versions, put formats and encryption keys.
 * @param blob the array of blob records for the messages.
 * @param userMetadata the array of userMetadata for the messages.
 * @param storeKeys the array of store keys for the messages.
 * @param encryptionKeys the array of encryption keys for the messages.
 * @param putFormats the array of Put Format class names to use to create the message streams.
 * @param headerVersions the array of Message Header versions to use for the messages.
 */
private void doSendWriteCompositeMessagesTest(byte[][] blob, byte[][] userMetadata, StoreKey[] storeKeys, ByteBuffer[] encryptionKeys, String[] putFormats, short[] headerVersions) throws MessageFormatException, IOException {
    String serviceIdPrefix = "serviceId";
    String ownerIdPrefix = "owner";
    String contentTypePrefix = "bin";
    short accountIdBase = 10;
    short containerIdBase = 2;
    BlobProperties[] properties = new BlobProperties[5];
    for (int i = 0; i < 5; i++) {
        properties[i] = new BlobProperties(blob[i].length, serviceIdPrefix + i, ownerIdPrefix + i, contentTypePrefix + i, false, 100, (short) (accountIdBase + i), (short) (containerIdBase + i), encryptionKeys[i] != null, null, null, null);
    }
    MessageFormatInputStream[] putStreams = new MessageFormatInputStream[5];
    for (int i = 0; i < 5; i++) {
        MessageFormatRecord.headerVersionToUse = headerVersions[i];
        if (putFormats[i].equals(PutMessageFormatInputStream.class.getSimpleName())) {
            putStreams[i] = new PutMessageFormatInputStream(storeKeys[i], (ByteBuffer) encryptionKeys[i].rewind(), properties[i], ByteBuffer.wrap(userMetadata[i]), new ByteBufferInputStream(ByteBuffer.wrap(blob[i])), blob[i].length, BlobType.DataBlob);
        } else {
            putStreams[i] = new PutMessageFormatBlobV1InputStream(storeKeys[i], properties[i], ByteBuffer.wrap(userMetadata[i]), new ByteBufferInputStream(ByteBuffer.wrap(blob[i])), blob[i].length, BlobType.DataBlob);
        }
    }
    int totalStreamSize = (int) Arrays.stream(putStreams).mapToLong(MessageFormatInputStream::getSize).sum();
    ByteBuffer compositeBuf = ByteBuffer.allocate(totalStreamSize);
    ArrayList<ByteBuffer> listbuf = new ArrayList<>();
    for (int i = 0; i < 5; i++) {
        ByteBuffer buf = ByteBuffer.allocate((int) putStreams[i].getSize());
        putStreams[i].read(buf.array());
        compositeBuf.put(buf.array());
        listbuf.add(buf);
    }
    MessageReadSet readSet = new MockMessageReadSet(listbuf, new ArrayList<>(Arrays.asList(storeKeys)));
    MetricRegistry registry = new MetricRegistry();
    MessageFormatMetrics metrics = new MessageFormatMetrics(registry);
    // get all
    MessageFormatSend send = new MessageFormatSend(readSet, MessageFormatFlags.All, metrics, new MockIdFactory());
    Assert.assertEquals(send.sizeInBytes(), totalStreamSize);
    Assert.assertEquals(5, send.getMessageMetadataList().size());
    for (int i = 0; i < 5; i++) {
        Assert.assertEquals(null, send.getMessageMetadataList().get(i));
    }
    ByteBuffer bufresult = ByteBuffer.allocate(totalStreamSize);
    WritableByteChannel channel = Channels.newChannel(new ByteBufferOutputStream(bufresult));
    while (!send.isSendComplete()) {
        send.writeTo(channel);
    }
    Assert.assertArrayEquals(compositeBuf.array(), bufresult.array());
    // get blob
    send = new MessageFormatSend(readSet, MessageFormatFlags.Blob, metrics, new MockIdFactory());
    int[] blobRecordSizes = new int[5];
    for (int i = 0; i < 5; i++) {
        blobRecordSizes[i] = (int) (putFormats[i].equals(PutMessageFormatInputStream.class.getSimpleName()) ? MessageFormatRecord.Blob_Format_V2.getBlobRecordSize(blob[i].length) : MessageFormatRecord.Blob_Format_V1.getBlobRecordSize(blob[i].length));
    }
    Assert.assertEquals(send.sizeInBytes(), (long) Arrays.stream(blobRecordSizes).sum());
    bufresult.clear();
    channel = Channels.newChannel(new ByteBufferOutputStream(bufresult));
    while (!send.isSendComplete()) {
        send.writeTo(channel);
    }
    int startOffset = 0;
    for (int i = 0; i < 5; i++) {
        DeserializedBlob deserializedBlob = MessageFormatRecord.deserializeAndGetBlobWithVersion(new ByteArrayInputStream(bufresult.array(), startOffset, blobRecordSizes[i]));
        Assert.assertEquals(putFormats[i].equals(PutMessageFormatInputStream.class.getSimpleName()) ? MessageFormatRecord.Blob_Version_V2 : MessageFormatRecord.Blob_Version_V1, deserializedBlob.getVersion());
        Assert.assertEquals(BlobType.DataBlob, deserializedBlob.getBlobData().getBlobType());
        Assert.assertEquals(blob[i].length, deserializedBlob.getBlobData().getSize());
        byte[] readBlob = new byte[blob[i].length];
        deserializedBlob.getBlobData().content().readBytes(readBlob);
        Assert.assertArrayEquals(blob[i], readBlob);
        deserializedBlob.getBlobData().release();
        if (headerVersions[i] == MessageFormatRecord.Message_Header_Version_V1) {
            Assert.assertEquals(null, send.getMessageMetadataList().get(i));
        } else {
            Assert.assertEquals(encryptionKeys[i].rewind(), send.getMessageMetadataList().get(i).getEncryptionKey());
        }
        startOffset += blobRecordSizes[i];
    }
    // get user metadata
    send = new MessageFormatSend(readSet, MessageFormatFlags.BlobUserMetadata, metrics, new MockIdFactory());
    int[] userMetadataSizes = new int[5];
    for (int i = 0; i < 5; i++) {
        userMetadataSizes[i] = MessageFormatRecord.UserMetadata_Format_V1.getUserMetadataSize(ByteBuffer.wrap(userMetadata[i]));
    }
    Assert.assertEquals(send.sizeInBytes(), (long) Arrays.stream(userMetadataSizes).sum());
    bufresult.clear();
    channel = Channels.newChannel(new ByteBufferOutputStream(bufresult));
    while (!send.isSendComplete()) {
        send.writeTo(channel);
    }
    startOffset = 0;
    for (int i = 0; i < 5; i++) {
        DeserializedUserMetadata deserializedUserMetadata = MessageFormatRecord.deserializeAndGetUserMetadataWithVersion(new ByteArrayInputStream(bufresult.array(), startOffset, userMetadataSizes[i]));
        Assert.assertEquals(MessageFormatRecord.UserMetadata_Version_V1, deserializedUserMetadata.getVersion());
        verifyBlobUserMetadata(userMetadata[i], deserializedUserMetadata.getUserMetadata());
        if (headerVersions[i] == MessageFormatRecord.Message_Header_Version_V1) {
            Assert.assertEquals(null, send.getMessageMetadataList().get(i));
        } else {
            Assert.assertEquals(encryptionKeys[i].rewind(), send.getMessageMetadataList().get(i).getEncryptionKey());
        }
        startOffset += userMetadataSizes[i];
    }
    // get blob properties
    send = new MessageFormatSend(readSet, MessageFormatFlags.BlobProperties, metrics, new MockIdFactory());
    int[] blobPropertiesSizes = new int[5];
    for (int i = 0; i < 5; i++) {
        blobPropertiesSizes[i] = MessageFormatRecord.BlobProperties_Format_V1.getBlobPropertiesRecordSize(properties[i]);
    }
    Assert.assertEquals(send.sizeInBytes(), (long) Arrays.stream(blobPropertiesSizes).sum());
    bufresult.clear();
    channel = Channels.newChannel(new ByteBufferOutputStream(bufresult));
    while (!send.isSendComplete()) {
        send.writeTo(channel);
    }
    startOffset = 0;
    for (int i = 0; i < 5; i++) {
        DeserializedBlobProperties deserializedBlobProperties = MessageFormatRecord.deserializeAndGetBlobPropertiesWithVersion(new ByteArrayInputStream(bufresult.array(), startOffset, blobPropertiesSizes[i]));
        Assert.assertEquals(MessageFormatRecord.BlobProperties_Version_V1, deserializedBlobProperties.getVersion());
        verifyBlobProperties(properties[i], deserializedBlobProperties.getBlobProperties());
        Assert.assertEquals(null, send.getMessageMetadataList().get(i));
        startOffset += blobPropertiesSizes[i];
    }
    // get blob info
    send = new MessageFormatSend(readSet, MessageFormatFlags.BlobInfo, metrics, new MockIdFactory());
    int[] blobInfoSizes = new int[5];
    for (int i = 0; i < 5; i++) {
        blobInfoSizes[i] = MessageFormatRecord.BlobProperties_Format_V1.getBlobPropertiesRecordSize(properties[i]) + MessageFormatRecord.UserMetadata_Format_V1.getUserMetadataSize(ByteBuffer.wrap(userMetadata[i]));
    }
    Assert.assertEquals(send.sizeInBytes(), (long) Arrays.stream(blobInfoSizes).sum());
    bufresult.clear();
    channel = Channels.newChannel(new ByteBufferOutputStream(bufresult));
    while (!send.isSendComplete()) {
        send.writeTo(channel);
    }
    startOffset = 0;
    for (int i = 0; i < 5; i++) {
        ByteArrayInputStream inputStream = new ByteArrayInputStream(bufresult.array(), startOffset, blobInfoSizes[i]);
        DeserializedBlobProperties deserializedBlobProperties = MessageFormatRecord.deserializeAndGetBlobPropertiesWithVersion(inputStream);
        DeserializedUserMetadata deserializedUserMetadata = MessageFormatRecord.deserializeAndGetUserMetadataWithVersion(inputStream);
        Assert.assertEquals(MessageFormatRecord.BlobProperties_Version_V1, deserializedBlobProperties.getVersion());
        verifyBlobProperties(properties[i], deserializedBlobProperties.getBlobProperties());
        Assert.assertEquals(MessageFormatRecord.UserMetadata_Version_V1, deserializedUserMetadata.getVersion());
        verifyBlobUserMetadata(userMetadata[i], deserializedUserMetadata.getUserMetadata());
        if (headerVersions[i] == MessageFormatRecord.Message_Header_Version_V1) {
            Assert.assertEquals(null, send.getMessageMetadataList().get(i));
        } else {
            Assert.assertEquals(encryptionKeys[i].rewind(), send.getMessageMetadataList().get(i).getEncryptionKey());
        }
        startOffset += blobInfoSizes[i];
    }
}