BlockDataManagerTest.java

/*
 * Licensed to the University of California, Berkeley under one or more contributor license
 * agreements. See the NOTICE file distributed with this work for additional information regarding
 * copyright ownership. The ASF licenses this file to You 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 tachyon.worker.block;


import java.io.IOException;
import java.util.LinkedList;
import java.util.Random;

import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import org.mockito.Mockito;

import tachyon.Constants;
import tachyon.Sessions;
import tachyon.client.BlockMasterClient;
import tachyon.client.FileSystemMasterClient;
import tachyon.conf.TachyonConf;
import tachyon.test.Tester;
import tachyon.thrift.FileInfo;
import tachyon.underfs.UnderFileSystem;
import tachyon.util.io.PathUtils;
import tachyon.worker.WorkerSource;
import tachyon.worker.block.meta.BlockMeta;
import tachyon.worker.block.meta.StorageDir;
import tachyon.worker.block.meta.TempBlockMeta;

public class BlockDataManagerTest implements Tester<BlockDataManager> {
  TestHarness mHarness;
  BlockDataManager.PrivateAccess mPrivateAccess;

  public void receiveAccess(Object access) {
    mPrivateAccess = (BlockDataManager.PrivateAccess) access;
  }

  class TestHarness {
    BlockMasterClient mBlockMasterClient;
    BlockStore mBlockStore;
    FileSystemMasterClient mFileSystemMasterClient;
    BlockHeartbeatReporter mHeartbeatReporter;
    BlockDataManager mManager;
    BlockMetricsReporter mMetricsReporter;
    Random mRandom;
    Sessions mSessions;
    TachyonConf mTachyonConf;
    UnderFileSystem mUfs;
    long mWorkerId;
    WorkerSource mWorkerSource;

    public TestHarness() throws IOException {
      mRandom = new Random();

      mBlockMasterClient = Mockito.mock(BlockMasterClient.class);
      mBlockStore = Mockito.mock(BlockStore.class);
      mFileSystemMasterClient = Mockito.mock(FileSystemMasterClient.class);
      mHeartbeatReporter = Mockito.mock(BlockHeartbeatReporter.class);
      mMetricsReporter = Mockito.mock(BlockMetricsReporter.class);
      mSessions = Mockito.mock(Sessions.class);
      mTachyonConf = Mockito.mock(TachyonConf.class);
      mUfs = Mockito.mock(UnderFileSystem.class);
      mWorkerId = mRandom.nextLong();
      mWorkerSource = Mockito.mock(WorkerSource.class);

      mManager = new BlockDataManager(mWorkerSource, mBlockMasterClient, mFileSystemMasterClient);
      mManager.setSessions(mSessions);
      mManager.setWorkerId(mWorkerId);

      mManager.grantAccess(BlockDataManagerTest.this); // initializes mPrivateAccess
      mPrivateAccess.setBlockStore(mBlockStore);
      mPrivateAccess.setHeartbeatReporter(mHeartbeatReporter);
      mPrivateAccess.setMetricsReporter(mMetricsReporter);
      mPrivateAccess.setTachyonConf(mTachyonConf);
      mPrivateAccess.setUnderFileSystem(mUfs);
    }
  }

  @Before
  public void initialize() throws IOException {
    mHarness = new TestHarness();
  }

  @Test
  public void abortBlockTest() throws Exception {
    long blockId = mHarness.mRandom.nextLong();
    long sessionId = mHarness.mRandom.nextLong();
    mHarness.mManager.abortBlock(sessionId, blockId);
    Mockito.verify(mHarness.mBlockStore).abortBlock(sessionId, blockId);
  }

  @Test
  public void accessBlockTest() throws Exception {
    long blockId = mHarness.mRandom.nextLong();
    long sessionId = mHarness.mRandom.nextLong();
    mHarness.mManager.accessBlock(sessionId, blockId);
    Mockito.verify(mHarness.mBlockStore).accessBlock(sessionId, blockId);
  }

  @Test
  public void addCheckpointTest() throws Exception {
    long fileId = mHarness.mRandom.nextLong();
    long fileSize = mHarness.mRandom.nextLong();
    long sessionId = mHarness.mRandom.nextLong();
    FileInfo fileInfo = new FileInfo();
    fileInfo.setPath("/foo/bar");
    String srcPath = "/tmp/" + fileId;
    String parentPath = "/tmp/foo";
    String dstPath = "/tmp/foo/bar";

    // TODO(jsimsa): Add test cases for error cases.
    Mockito.when(mHarness.mTachyonConf
        .get(Constants.UNDERFS_DATA_FOLDER, Constants.DEFAULT_DATA_FOLDER)).thenReturn("/tmp");
    Mockito.when(mHarness.mSessions.getSessionUfsTempFolder(sessionId)).thenReturn("/tmp");
    Mockito.when(mHarness.mFileSystemMasterClient.getFileInfo(fileId)).thenReturn(fileInfo);
    Mockito.when(mHarness.mUfs.exists(parentPath)).thenReturn(true);
    Mockito.when(mHarness.mUfs.mkdirs(parentPath, true)).thenReturn(true);
    Mockito.when(mHarness.mUfs.rename(srcPath, dstPath)).thenReturn(true);
    Mockito.when(mHarness.mUfs.getFileSize(dstPath)).thenReturn(fileSize);
    mHarness.mManager.addCheckpoint(sessionId, fileId);
    Mockito.verify(mHarness.mFileSystemMasterClient)
        .addCheckpoint(mHarness.mWorkerId, fileId, fileSize, dstPath);
  }

  @Test
  public void cleanupSessionsTest() throws Exception {
    long sessionId = 1;
    LinkedList<Long> sessions = new LinkedList<Long>();
    sessions.add(sessionId);

    Mockito.when(mHarness.mSessions.getTimedOutSessions()).thenReturn(sessions);
    mHarness.mManager.cleanupSessions();
    Mockito.verify(mHarness.mSessions).removeSession(sessionId);
    Mockito.verify(mHarness.mBlockStore).cleanupSession(sessionId);
  }

  @Test
  public void commitBlockTest() throws Exception {
    long blockId = mHarness.mRandom.nextLong();
    long length = mHarness.mRandom.nextLong();
    long lockId = mHarness.mRandom.nextLong();
    long sessionId = mHarness.mRandom.nextLong();
    long usedBytes = mHarness.mRandom.nextLong();
    int tierAlias = 1;
    LinkedList<Long> usedBytesOnTiers = new LinkedList<Long>();
    usedBytesOnTiers.add(usedBytes);
    BlockMeta blockMeta = Mockito.mock(BlockMeta.class);
    BlockStoreLocation blockStoreLocation = Mockito.mock(BlockStoreLocation.class);
    BlockStoreMeta blockStoreMeta = Mockito.mock(BlockStoreMeta.class);

    Mockito.when(mHarness.mBlockStore.lockBlock(sessionId, blockId)).thenReturn(lockId);
    Mockito.when(mHarness.mBlockStore.getBlockMeta(sessionId, blockId, lockId))
        .thenReturn(blockMeta);
    Mockito.when(mHarness.mBlockStore.getBlockStoreMeta()).thenReturn(blockStoreMeta);
    Mockito.when(blockMeta.getBlockLocation()).thenReturn(blockStoreLocation);
    Mockito.when(blockStoreLocation.tierAlias()).thenReturn(tierAlias);
    Mockito.when(blockMeta.getBlockSize()).thenReturn(length);
    Mockito.when(blockStoreMeta.getUsedBytesOnTiers()).thenReturn(usedBytesOnTiers);

    mHarness.mManager.commitBlock(sessionId, blockId);
    Mockito.verify(mHarness.mBlockMasterClient)
        .workerCommitBlock(mHarness.mWorkerId, usedBytes, tierAlias, blockId, length);
    Mockito.verify(mHarness.mBlockStore).unlockBlock(lockId);
  }

  @Test
  public void CreateBlockTest() throws Exception {
    long blockId = mHarness.mRandom.nextLong();
    long initialBytes = mHarness.mRandom.nextLong();
    long sessionId = mHarness.mRandom.nextLong();
    int tierAlias = 1;
    BlockStoreLocation location = BlockStoreLocation.anyDirInTier(tierAlias);
    StorageDir storageDir = Mockito.mock(StorageDir.class);
    TempBlockMeta meta = new TempBlockMeta(sessionId, blockId, initialBytes, storageDir);

    Mockito.when(mHarness.mBlockStore
        .createBlockMeta(sessionId, blockId, location, initialBytes)).thenReturn(meta);
    Mockito.when(storageDir.getDirPath()).thenReturn("/tmp");
    Assert.assertEquals(PathUtils.concatPath("/tmp", sessionId, blockId),
        mHarness.mManager.createBlock(sessionId, blockId, tierAlias, initialBytes));
  }

  // TODO(jsimsa): Write unit tests for untested public methods.
}