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

Example 21 with MockId

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

the class MessageFormatInputStreamTest method messageFormatPutRecordsTest.

private void messageFormatPutRecordsTest(short blobVersion, BlobType blobType, short headerVersion) throws IOException, MessageFormatException {
    StoreKey key = new MockId("id1");
    StoreKeyFactory keyFactory = new MockIdFactory();
    short accountId = Utils.getRandomShort(TestUtils.RANDOM);
    short containerId = Utils.getRandomShort(TestUtils.RANDOM);
    BlobProperties prop = new BlobProperties(10, "servid", accountId, containerId, false);
    byte[] encryptionKey = new byte[100];
    new Random().nextBytes(encryptionKey);
    byte[] usermetadata = new byte[1000];
    new Random().nextBytes(usermetadata);
    int blobContentSize = 2000;
    byte[] data = new byte[blobContentSize];
    new Random().nextBytes(data);
    short lifeVersion = 1;
    long blobSize = -1;
    MessageFormatRecord.headerVersionToUse = headerVersion;
    if (blobVersion == MessageFormatRecord.Blob_Version_V1) {
        blobSize = MessageFormatRecord.Blob_Format_V1.getBlobRecordSize(blobContentSize);
    } else if (blobVersion == MessageFormatRecord.Blob_Version_V2 && blobType == BlobType.DataBlob) {
        blobSize = (int) MessageFormatRecord.Blob_Format_V2.getBlobRecordSize(blobContentSize);
    } else if (blobVersion == MessageFormatRecord.Blob_Version_V2 && blobType == BlobType.MetadataBlob) {
        ByteBuffer byteBufferBlob = MessageFormatTestUtils.getBlobContentForMetadataBlob(blobContentSize);
        data = byteBufferBlob.array();
        blobContentSize = data.length;
        blobSize = (int) MessageFormatRecord.Blob_Format_V2.getBlobRecordSize(blobContentSize);
    }
    ByteBufferInputStream stream = new ByteBufferInputStream(ByteBuffer.wrap(data));
    MessageFormatInputStream messageFormatStream = (blobVersion == MessageFormatRecord.Blob_Version_V2) ? new PutMessageFormatInputStream(key, ByteBuffer.wrap(encryptionKey), prop, ByteBuffer.wrap(usermetadata), stream, blobContentSize, blobType, lifeVersion) : new PutMessageFormatBlobV1InputStream(key, prop, ByteBuffer.wrap(usermetadata), stream, blobContentSize, blobType);
    int headerSize = MessageFormatRecord.getHeaderSizeForVersion(headerVersion);
    int blobEncryptionKeySize = headerVersion != MessageFormatRecord.Message_Header_Version_V1 ? MessageFormatRecord.BlobEncryptionKey_Format_V1.getBlobEncryptionKeyRecordSize(ByteBuffer.wrap(encryptionKey)) : 0;
    int blobPropertiesRecordSize = MessageFormatRecord.BlobProperties_Format_V1.getBlobPropertiesRecordSize(prop);
    int userMetadataSize = MessageFormatRecord.UserMetadata_Format_V1.getUserMetadataSize(ByteBuffer.wrap(usermetadata));
    Assert.assertEquals(messageFormatStream.getSize(), headerSize + blobEncryptionKeySize + blobPropertiesRecordSize + userMetadataSize + blobSize + key.sizeInBytes());
    // verify header
    byte[] headerOutput = new byte[headerSize];
    messageFormatStream.read(headerOutput);
    ByteBuffer headerBuf = ByteBuffer.wrap(headerOutput);
    Assert.assertEquals(headerVersion, headerBuf.getShort());
    if (headerVersion == MessageFormatRecord.Message_Header_Version_V3) {
        Assert.assertEquals(lifeVersion, headerBuf.getShort());
    }
    Assert.assertEquals(blobEncryptionKeySize + blobPropertiesRecordSize + userMetadataSize + blobSize, headerBuf.getLong());
    switch(headerVersion) {
        case MessageFormatRecord.Message_Header_Version_V1:
            Assert.assertEquals(headerSize + key.sizeInBytes(), headerBuf.getInt());
            Assert.assertEquals(MessageFormatRecord.Message_Header_Invalid_Relative_Offset, headerBuf.getInt());
            Assert.assertEquals(headerSize + key.sizeInBytes() + blobPropertiesRecordSize, headerBuf.getInt());
            Assert.assertEquals(headerSize + key.sizeInBytes() + blobPropertiesRecordSize + userMetadataSize, headerBuf.getInt());
            break;
        default:
            // case MessageFormatRecord.Message_Header_Version_V2 or V3:
            Assert.assertEquals(headerSize + key.sizeInBytes(), headerBuf.getInt());
            Assert.assertEquals(headerSize + key.sizeInBytes() + blobEncryptionKeySize, headerBuf.getInt());
            Assert.assertEquals(MessageFormatRecord.Message_Header_Invalid_Relative_Offset, headerBuf.getInt());
            Assert.assertEquals(headerSize + key.sizeInBytes() + blobEncryptionKeySize + blobPropertiesRecordSize, headerBuf.getInt());
            Assert.assertEquals(headerSize + key.sizeInBytes() + blobEncryptionKeySize + blobPropertiesRecordSize + userMetadataSize, headerBuf.getInt());
    }
    Crc32 crc = new Crc32();
    crc.update(headerOutput, 0, headerSize - MessageFormatRecord.Crc_Size);
    Assert.assertEquals(crc.getValue(), headerBuf.getLong());
    // verify handle
    byte[] handleOutput = new byte[key.sizeInBytes()];
    ByteBuffer handleOutputBuf = ByteBuffer.wrap(handleOutput);
    messageFormatStream.read(handleOutput);
    byte[] dest = new byte[key.sizeInBytes()];
    handleOutputBuf.get(dest);
    Assert.assertArrayEquals(dest, key.toBytes());
    // verify encryption key
    if (headerVersion != MessageFormatRecord.Message_Header_Version_V1) {
        byte[] blobEncryptionKeyOutput = new byte[blobEncryptionKeySize];
        ByteBuffer blobEncryptionKeyBuf = ByteBuffer.wrap(blobEncryptionKeyOutput);
        messageFormatStream.read(blobEncryptionKeyOutput);
        Assert.assertEquals(blobEncryptionKeyBuf.getShort(), MessageFormatRecord.Blob_Encryption_Key_V1);
        Assert.assertEquals(blobEncryptionKeyBuf.getInt(), 100);
        dest = new byte[100];
        blobEncryptionKeyBuf.get(dest);
        Assert.assertArrayEquals(dest, encryptionKey);
        crc = new Crc32();
        crc.update(blobEncryptionKeyOutput, 0, blobEncryptionKeySize - MessageFormatRecord.Crc_Size);
        Assert.assertEquals(crc.getValue(), blobEncryptionKeyBuf.getLong());
    }
    // verify blob properties
    byte[] blobPropertiesOutput = new byte[blobPropertiesRecordSize];
    ByteBuffer blobPropertiesBuf = ByteBuffer.wrap(blobPropertiesOutput);
    messageFormatStream.read(blobPropertiesOutput);
    Assert.assertEquals(blobPropertiesBuf.getShort(), 1);
    BlobProperties propOutput = BlobPropertiesSerDe.getBlobPropertiesFromStream(new DataInputStream(new ByteBufferInputStream(blobPropertiesBuf)));
    Assert.assertEquals(10, propOutput.getBlobSize());
    Assert.assertEquals("servid", propOutput.getServiceId());
    Assert.assertEquals("AccountId mismatch", accountId, propOutput.getAccountId());
    Assert.assertEquals("ContainerId mismatch", containerId, propOutput.getContainerId());
    crc = new Crc32();
    crc.update(blobPropertiesOutput, 0, blobPropertiesRecordSize - MessageFormatRecord.Crc_Size);
    Assert.assertEquals(crc.getValue(), blobPropertiesBuf.getLong());
    // verify user metadata
    byte[] userMetadataOutput = new byte[userMetadataSize];
    ByteBuffer userMetadataBuf = ByteBuffer.wrap(userMetadataOutput);
    messageFormatStream.read(userMetadataOutput);
    Assert.assertEquals(userMetadataBuf.getShort(), 1);
    Assert.assertEquals(userMetadataBuf.getInt(), 1000);
    dest = new byte[1000];
    userMetadataBuf.get(dest);
    Assert.assertArrayEquals(dest, usermetadata);
    crc = new Crc32();
    crc.update(userMetadataOutput, 0, userMetadataSize - MessageFormatRecord.Crc_Size);
    Assert.assertEquals(crc.getValue(), userMetadataBuf.getLong());
    // verify blob
    CrcInputStream crcstream = new CrcInputStream(messageFormatStream);
    DataInputStream streamData = new DataInputStream(crcstream);
    Assert.assertEquals(streamData.readShort(), blobVersion);
    if (blobVersion == MessageFormatRecord.Blob_Version_V2) {
        Assert.assertEquals(streamData.readShort(), blobType.ordinal());
    }
    Assert.assertEquals(streamData.readLong(), blobContentSize);
    for (int i = 0; i < blobContentSize; i++) {
        Assert.assertEquals((byte) streamData.read(), data[i]);
    }
    long crcVal = crcstream.getValue();
    Assert.assertEquals(crcVal, streamData.readLong());
    // Verify Blob All
    stream = new ByteBufferInputStream(ByteBuffer.wrap(data));
    messageFormatStream = (blobVersion == MessageFormatRecord.Blob_Version_V2) ? new PutMessageFormatInputStream(key, ByteBuffer.wrap(encryptionKey), prop, ByteBuffer.wrap(usermetadata), stream, blobContentSize, blobType) : new PutMessageFormatBlobV1InputStream(key, prop, ByteBuffer.wrap(usermetadata), stream, blobContentSize, blobType);
    int totalSize;
    switch(headerVersion) {
        case MessageFormatRecord.Message_Header_Version_V1:
            totalSize = headerSize + key.sizeInBytes() + blobPropertiesRecordSize + userMetadataSize + (int) blobSize;
            break;
        default:
            // case MessageFormatRecord.Message_Header_Version_V2:
            totalSize = headerSize + key.sizeInBytes() + blobEncryptionKeySize + blobPropertiesRecordSize + userMetadataSize + (int) blobSize;
    }
    ByteBuffer allBuf = ByteBuffer.allocate(totalSize);
    messageFormatStream.read(allBuf.array());
    BlobAll blobAll = MessageFormatRecord.deserializeBlobAll(new ByteBufferInputStream(allBuf), keyFactory);
    Assert.assertEquals(key, blobAll.getStoreKey());
    Assert.assertArrayEquals(usermetadata, blobAll.getBlobInfo().getUserMetadata());
    Assert.assertEquals(blobContentSize, blobAll.getBlobData().getSize());
    Assert.assertEquals(blobType, blobAll.getBlobData().getBlobType());
    if (headerVersion != MessageFormatRecord.Message_Header_Version_V1) {
        Assert.assertEquals(ByteBuffer.wrap(encryptionKey), blobAll.getBlobEncryptionKey());
    } else {
        Assert.assertEquals(null, blobAll.getBlobEncryptionKey());
    }
    ByteBuf byteBuf = blobAll.getBlobData().content();
    try {
        Assert.assertEquals(Unpooled.wrappedBuffer(data), byteBuf);
    } finally {
        byteBuf.release();
    }
}