com.linkedin.pinot.controller.validation.StorageQuotaCheckerTest.java Source code

Introduction

Here is the source code for com.linkedin.pinot.controller.validation.StorageQuotaCheckerTest.java

Source

/**
 * Copyright (C) 2014-2016 LinkedIn Corp. (pinot-core@linkedin.com)
 *
 * 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 com.linkedin.pinot.controller.validation;

import com.linkedin.pinot.common.config.AbstractTableConfig;
import com.linkedin.pinot.common.config.QuotaConfig;
import com.linkedin.pinot.common.config.SegmentsValidationAndRetentionConfig;
import com.linkedin.pinot.controller.api.restlet.resources.TableSizeReader;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.Arrays;
import org.apache.commons.io.FileUtils;
import org.mockito.invocation.InvocationOnMock;
import org.mockito.stubbing.Answer;
import org.testng.Assert;
import org.testng.annotations.AfterClass;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.Test;

import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

public class StorageQuotaCheckerTest {
    TableSizeReader tableSizeReader;
    AbstractTableConfig tableConfig;
    private QuotaConfig quotaConfig;
    private SegmentsValidationAndRetentionConfig validationConfig;
    private final File TEST_DIR = new File(StorageQuotaCheckerTest.class.getName());

    @BeforeClass
    public void setUp() {
        tableSizeReader = mock(TableSizeReader.class);
        tableConfig = mock(AbstractTableConfig.class);
        quotaConfig = mock(QuotaConfig.class);
        validationConfig = mock(SegmentsValidationAndRetentionConfig.class);
        when(tableConfig.getValidationConfig()).thenReturn(validationConfig);
        when(validationConfig.getReplicationNumber()).thenReturn(2);
        TEST_DIR.mkdirs();
    }

    @AfterClass
    public void tearDown() {
        FileUtils.deleteQuietly(TEST_DIR);
    }

    @Test
    public void testNoQuota() {
        StorageQuotaChecker checker = new StorageQuotaChecker(tableConfig, tableSizeReader);
        when(tableConfig.getQuotaConfig()).thenReturn(null);
        StorageQuotaChecker.QuotaCheckerResponse res = checker.isSegmentStorageWithinQuota(TEST_DIR, "myTable",
                "segment", 1000);
        Assert.assertTrue(res.isSegmentWithinQuota);
    }

    @Test
    public void testNoStorageQuotaConfig() {
        StorageQuotaChecker checker = new StorageQuotaChecker(tableConfig, tableSizeReader);
        when(tableConfig.getQuotaConfig()).thenReturn(quotaConfig);
        when(quotaConfig.storageSizeBytes()).thenReturn(-1L);
        StorageQuotaChecker.QuotaCheckerResponse res = checker.isSegmentStorageWithinQuota(TEST_DIR, "myTable",
                "segment", 1000);
        Assert.assertTrue(res.isSegmentWithinQuota);
    }

    public void setupTableSegmentSize(final long tableSize, final long segmentSize) {
        when(tableSizeReader.getTableSubtypeSize("testTable", 1000))
                .thenAnswer(new Answer<TableSizeReader.TableSubTypeSizeDetails>() {
                    @Override
                    public TableSizeReader.TableSubTypeSizeDetails answer(InvocationOnMock invocationOnMock)
                            throws Throwable {
                        TableSizeReader.TableSubTypeSizeDetails sizeDetails = new TableSizeReader.TableSubTypeSizeDetails();
                        sizeDetails.estimatedSizeInBytes = tableSize;
                        TableSizeReader.SegmentSizeDetails segSizeDetails = new TableSizeReader.SegmentSizeDetails();
                        segSizeDetails.estimatedSizeInBytes = segmentSize;
                        sizeDetails.segments.put("segment1", segSizeDetails);
                        return sizeDetails;
                    }
                });
    }

    @Test
    public void testWithinQuota() throws IOException {
        File tempFile = new File(TEST_DIR, "small_file");
        tempFile.createNewFile();
        byte[] data = new byte[1024];
        Arrays.fill(data, (byte) 1);
        try (FileOutputStream ostr = new FileOutputStream(tempFile)) {
            ostr.write(data);
        }
        setupTableSegmentSize(5800, 900);
        when(tableConfig.getQuotaConfig()).thenReturn(quotaConfig);
        when(quotaConfig.storageSizeBytes()).thenReturn(3000L);
        StorageQuotaChecker checker = new StorageQuotaChecker(tableConfig, tableSizeReader);
        StorageQuotaChecker.QuotaCheckerResponse response = checker.isSegmentStorageWithinQuota(TEST_DIR,
                "testTable", "segment1", 1000);
        Assert.assertTrue(response.isSegmentWithinQuota);

        when(quotaConfig.storageSizeBytes()).thenReturn(2800L);
        response = checker.isSegmentStorageWithinQuota(TEST_DIR, "testTable", "segment1", 1000);
        Assert.assertFalse(response.isSegmentWithinQuota);
    }
}