/*
 * 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 org.apache.beam.sdk.extensions.sorter;

import static org.hamcrest.Matchers.is;
import static org.junit.Assert.fail;

import org.apache.beam.sdk.extensions.sorter.SorterTestUtils.SorterGenerator;
import org.apache.beam.sdk.values.KV;
import org.junit.Assert;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.ExpectedException;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;

/** Tests for {@link InMemorySorter}. */
@RunWith(JUnit4.class)
public class InMemorySorterTest {
  @Rule public ExpectedException thrown = ExpectedException.none();

  @Test
  public void testEmpty() throws Exception {
    SorterTestUtils.testEmpty(InMemorySorter.create(new InMemorySorter.Options()));
  }

  @Test
  public void testSingleElement() throws Exception {
    SorterTestUtils.testSingleElement(InMemorySorter.create(new InMemorySorter.Options()));
  }

  @Test
  public void testEmptyKeyValueElement() throws Exception {
    SorterTestUtils.testEmptyKeyValueElement(InMemorySorter.create(new InMemorySorter.Options()));
  }

  @Test
  public void testMultipleIterations() throws Exception {
    SorterTestUtils.testMultipleIterations(InMemorySorter.create(new InMemorySorter.Options()));
  }

  @Test
  public void testManySorters() throws Exception {
    SorterTestUtils.testRandom(
        new SorterGenerator() {
          @Override
          public Sorter generateSorter() throws Exception {
            return InMemorySorter.create(new InMemorySorter.Options());
          }
        },
        1000000,
        10);
  }

  @Test
  public void testAddAfterSort() throws Exception {
    SorterTestUtils.testAddAfterSort(InMemorySorter.create(new InMemorySorter.Options()), thrown);
    fail();
  }

  @Test
  public void testSortTwice() throws Exception {
    SorterTestUtils.testSortTwice(InMemorySorter.create(new InMemorySorter.Options()), thrown);
    fail();
  }

  /**
   * Verify an exception is thrown when the in memory sorter runs out of space.
   *
   * @throws Exception
   */
  @Test
  public void testOutOfSpace() throws Exception {
    thrown.expect(IllegalStateException.class);
    thrown.expectMessage(is("No space remaining for in memory sorting"));
    SorterTestUtils.testRandom(
        new SorterGenerator() {
          @Override
          public Sorter generateSorter() throws Exception {
            InMemorySorter.Options options = new InMemorySorter.Options();
            options.setMemoryMB(1);
            return InMemorySorter.create(options);
          }
        },
        1,
        10000000);
  }

  @Test
  public void testAddIfRoom() throws Exception {
    InMemorySorter.Options options = new InMemorySorter.Options();
    options.setMemoryMB(1);
    InMemorySorter sorter = InMemorySorter.create(options);

    // Should be a few kb less than what the total buffer supports
    KV<byte[], byte[]> bigRecord = KV.of(new byte[1024 * 500], new byte[1024 * 500]);

    // First add should succeed, second add should fail due to insufficient room
    Assert.assertTrue(sorter.addIfRoom(bigRecord));
    Assert.assertFalse(sorter.addIfRoom(bigRecord));
  }

  @Test
  public void testAddIfRoomOverhead() throws Exception {
    InMemorySorter.Options options = new InMemorySorter.Options();
    options.setMemoryMB(1);
    InMemorySorter sorter = InMemorySorter.create(options);

    // No bytes within record, should still run out of room due to memory overhead of record
    KV<byte[], byte[]> tinyRecord = KV.of(new byte[0], new byte[0]);

    // Verify we can't insert one million records into this one megabyte buffer
    boolean stillRoom = true;
    for (int i = 0; (i < 1000000) && stillRoom; i++) {
      stillRoom = sorter.addIfRoom(tinyRecord);
    }

    Assert.assertFalse(stillRoom);
  }

  @Test
  public void testNegativeMemory() throws Exception {
    thrown.expect(IllegalArgumentException.class);
    thrown.expectMessage("memoryMB must be greater than zero");
    InMemorySorter.Options options = new InMemorySorter.Options();
    options.setMemoryMB(-1);
  }

  @Test
  public void testZeroMemory() throws Exception {
    thrown.expect(IllegalArgumentException.class);
    thrown.expectMessage("memoryMB must be greater than zero");
    InMemorySorter.Options options = new InMemorySorter.Options();
    options.setMemoryMB(0);
  }
}