TimePartitionedStreamTestBase.java

/*
 * Copyright © 2014-2015 Cask Data, Inc.
 *
 * 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 co.cask.cdap.data.stream;

import co.cask.cdap.api.flow.flowlet.StreamEvent;
import co.cask.cdap.common.io.Locations;
import com.google.common.collect.Lists;
import org.apache.twill.filesystem.Location;
import org.apache.twill.filesystem.LocationFactory;
import org.junit.Assert;
import org.junit.ClassRule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;

import java.io.IOException;
import java.util.Collections;
import java.util.List;

/**
 * Tests for {@link co.cask.cdap.data.stream.TimePartitionedStreamFileWriter}.
 */
public abstract class TimePartitionedStreamTestBase {

  @ClassRule
  public static TemporaryFolder tmpFolder = new TemporaryFolder();

  protected abstract LocationFactory getLocationFactory();

  @Test
  public void testTimePartition() throws IOException {
    // Create time partition file of 1 seconds each.
    String streamName = "stream";
    Location streamLocation = getLocationFactory().create(streamName);
    streamLocation.mkdirs();
    TimePartitionedStreamFileWriter writer = new TimePartitionedStreamFileWriter(streamLocation, 1000, "file", 100);

    // Write 2 events per millis for 3 seconds, starting at 0.5 second.
    long timeBase = 500;
    for (int i = 0; i < 3; i++) {
      for (int j = 0; j < 1000; j++) {
        long offset = i * 1000 + j;
        long timestamp = timeBase + offset;
        writer.append(StreamFileTestUtils.createEvent(timestamp, "Testing " + offset + " 0"));
        writer.append(StreamFileTestUtils.createEvent(timestamp, "Testing " + offset + " 1"));
      }
    }

    writer.close();

    // There should be four partition directory (500-1000, 1000-2000, 2000-3000, 3000-3500).
    List<Location> partitionDirs = Lists.newArrayList(streamLocation.list());
    Assert.assertEquals(4, partitionDirs.size());

    // The start time for the partitions should be 0, 1000, 2000, 3000
    Collections.sort(partitionDirs, Locations.LOCATION_COMPARATOR);
    for (int i = 0; i < 4; i++) {
      Assert.assertEquals(i * 1000, StreamUtils.getPartitionStartTime(partitionDirs.get(i).getName()));
    }
  }

  @Test
  public void testAppendAll() throws IOException {
    // Create time partition file of 1 seconds each.
    String streamName = "testAppendAll";
    Location streamLocation = getLocationFactory().create(streamName);
    streamLocation.mkdirs();
    TimePartitionedStreamFileWriter writer = new TimePartitionedStreamFileWriter(streamLocation, 1000, "file", 100);

    // Write 2 events per millis for 3 seconds, starting at 0.5 second.
    List<StreamEvent> events = Lists.newArrayList();
    long timeBase = 500;
    for (int i = 0; i < 3; i++) {
      for (int j = 0; j < 1000; j++) {
        long offset = i * 1000 + j;
        long timestamp = timeBase + offset;
        events.add(StreamFileTestUtils.createEvent(timestamp, "Testing " + offset + " 0"));
        events.add(StreamFileTestUtils.createEvent(timestamp, "Testing " + offset + " 1"));
      }
    }
    writer.appendAll(events.iterator());
    writer.close();

    // There should be four partition directory (500-1000, 1000-2000, 2000-3000, 3000-3500).
    List<Location> partitionDirs = Lists.newArrayList(streamLocation.list());
    Assert.assertEquals(4, partitionDirs.size());

    // The start time for the partitions should be 0, 1000, 2000, 3000
    Collections.sort(partitionDirs, Locations.LOCATION_COMPARATOR);
    for (int i = 0; i < 4; i++) {
      Assert.assertEquals(i * 1000, StreamUtils.getPartitionStartTime(partitionDirs.get(i).getName()));
    }
  }
}