gobblin.data.management.copy.predicates.RegistrationTimeSkipPredicateTest.java Source code

Introduction

Here is the source code for gobblin.data.management.copy.predicates.RegistrationTimeSkipPredicateTest.java

Source

/*
 * Licensed to the Apache Software Foundation (ASF) 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 gobblin.data.management.copy.predicates;

import java.util.Map;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileStatus;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.hive.ql.metadata.Partition;
import org.mockito.Mockito;
import org.testng.Assert;
import org.testng.annotations.Test;

import com.google.common.base.Optional;
import com.google.common.collect.Maps;

import gobblin.data.management.copy.CopyConfiguration;
import gobblin.data.management.copy.CopyContext;
import gobblin.data.management.copy.hive.HiveCopyEntityHelper;
import gobblin.data.management.copy.hive.HiveDataset;
import gobblin.data.management.copy.hive.HivePartitionFileSet;

public class RegistrationTimeSkipPredicateTest {

    @Test
    public void test() throws Exception {

        Path partition1Path = new Path("/path/to/partition1");
        long modTime = 100000;

        CopyContext copyContext = new CopyContext();
        CopyConfiguration copyConfiguration = Mockito.mock(CopyConfiguration.class);
        Mockito.doReturn(copyContext).when(copyConfiguration).getCopyContext();
        HiveDataset dataset = Mockito.mock(HiveDataset.class);
        FileSystem fs = Mockito.spy(FileSystem.getLocal(new Configuration()));
        FileStatus status = new FileStatus(1, false, 1, 1, modTime, partition1Path);
        Path qualifiedPath = fs.makeQualified(partition1Path);
        Mockito.doReturn(status).when(fs).getFileStatus(qualifiedPath);
        Mockito.doReturn(status).when(fs).getFileStatus(partition1Path);
        Mockito.doReturn(fs).when(dataset).getFs();

        HiveCopyEntityHelper helper = Mockito.mock(HiveCopyEntityHelper.class);
        Mockito.doReturn(copyConfiguration).when(helper).getConfiguration();
        Mockito.doReturn(dataset).when(helper).getDataset();

        RegistrationTimeSkipPredicate predicate = new RegistrationTimeSkipPredicate(helper);

        // partition exists, but registration time before modtime => don't skip
        HivePartitionFileSet pc = createPartitionCopy(partition1Path, modTime - 1, true);
        Assert.assertFalse(predicate.apply(pc));

        // partition exists, registration time equal modtime => don't skip
        pc = createPartitionCopy(partition1Path, modTime, true);
        Assert.assertFalse(predicate.apply(pc));

        // partition exists, registration time larger modtime => do skip
        pc = createPartitionCopy(partition1Path, modTime + 1, true);
        Assert.assertTrue(predicate.apply(pc));

        // partition doesn't exist => don't skip
        pc = createPartitionCopy(partition1Path, modTime + 1, false);
        Assert.assertFalse(predicate.apply(pc));

        // partition exists but is not annotated => don't skip
        pc = createPartitionCopy(partition1Path, modTime + 1, true);
        pc.getExistingTargetPartition().get().getParameters().clear();
        Assert.assertFalse(predicate.apply(pc));

    }

    public HivePartitionFileSet createPartitionCopy(Path location, long registrationGenerationTime,
            boolean targetPartitionExists) {
        HivePartitionFileSet partitionCopy = Mockito.mock(HivePartitionFileSet.class);

        Partition partition = Mockito.mock(Partition.class);
        Mockito.doReturn(location).when(partition).getDataLocation();
        Mockito.doReturn(partition).when(partitionCopy).getPartition();

        if (targetPartitionExists) {

            Partition targetPartition = Mockito.mock(Partition.class);

            Map<String, String> parameters = Maps.newHashMap();
            parameters.put(HiveDataset.REGISTRATION_GENERATION_TIME_MILLIS,
                    Long.toString(registrationGenerationTime));
            Mockito.doReturn(parameters).when(targetPartition).getParameters();

            Mockito.doReturn(Optional.of(targetPartition)).when(partitionCopy).getExistingTargetPartition();
        } else {
            Mockito.doReturn(Optional.absent()).when(partitionCopy).getExistingTargetPartition();
        }

        return partitionCopy;
    }

}