<commit_before>package org.zalando.nakadi.service.subscription;

import org.hamcrest.MatcherAssert;
import org.junit.Assert;
import org.junit.Test;
import org.zalando.nakadi.service.subscription.model.Partition;
import org.zalando.nakadi.service.subscription.model.Partition.PartitionKey;
import org.zalando.nakadi.service.subscription.model.Session;

import java.util.stream.Stream;

import static org.hamcrest.Matchers.emptyArray;
import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
import static org.zalando.nakadi.service.subscription.model.Partition.State.ASSIGNED;
import static org.zalando.nakadi.service.subscription.model.Partition.State.REASSIGNING;

public class ExactWeightRebalancerTest {

    @Test(expected = IllegalArgumentException.class)
    public void splitByWeightShouldAcceptOnlyCorrectData1() {
        ExactWeightRebalancer.splitByWeight(1, new int[]{1, 1});
    }

    @Test(expected = IllegalArgumentException.class)
    public void splitByWeightShouldAcceptOnlyCorrectData2() {
        ExactWeightRebalancer.splitByWeight(2, new int[]{1, 0});
    }

    @Test
    public void splitByWeightMustCorrectlyWorkOnDifferentValues() {
        assertArrayEquals(
                new int[]{1, 1},
                ExactWeightRebalancer.splitByWeight(2, new int[]{1, 1}));
        assertArrayEquals(
                new int[]{2, 1},
                ExactWeightRebalancer.splitByWeight(3, new int[]{1, 1}));
        assertArrayEquals(
                new int[]{1, 2},
                ExactWeightRebalancer.splitByWeight(3, new int[]{1, 2}));
        assertArrayEquals(
                new int[]{34, 33, 33},
                ExactWeightRebalancer.splitByWeight(100, new int[]{1, 1, 1}));
        assertArrayEquals(
                new int[]{26, 25, 50},
                ExactWeightRebalancer.splitByWeight(101, new int[]{1, 1, 2}));
    }

    @Test
    public void rebalanceShouldHaveEmptyChangesetForBalancedData() {
        final ExactWeightRebalancer rebalancer = new ExactWeightRebalancer();

        // 1. Data contains only assigned
        final Session[] sessions = new Session[]{
                new Session("0", 1),
                new Session("1", 1)};
        MatcherAssert.assertThat(rebalancer.apply(sessions,
                new Partition[] {
                        new Partition(new PartitionKey("0", "0"), "0", null, ASSIGNED),
                        new Partition(new PartitionKey("0", "1"), "1", null, ASSIGNED),
                        new Partition(new PartitionKey("1", "0"), "1", null, ASSIGNED),
                        new Partition(new PartitionKey("1", "1"), "0", null, ASSIGNED)}),
                emptyArray());

        // 2. Data contains reassinging
        MatcherAssert.assertThat(rebalancer.apply(sessions,
                new Partition[] {
                        new Partition(new PartitionKey("0", "0"), "0", null, ASSIGNED),
                        new Partition(new PartitionKey("0", "1"), "0", "1", REASSIGNING),
                        new Partition(new PartitionKey("1", "0"), "1", null, ASSIGNED),
                        new Partition(new PartitionKey("1", "1"), "0", null, ASSIGNED)}),
                emptyArray());

        // 3. Data contains only reassinging
        MatcherAssert.assertThat(rebalancer.apply(sessions,
                new Partition[] {
                        new Partition(new PartitionKey("0", "0"), "0", "1", REASSIGNING),
                        new Partition(new PartitionKey("0", "1"), "0", "1", REASSIGNING),
                        new Partition(new PartitionKey("1", "0"), "1", "0", REASSIGNING),
                        new Partition(new PartitionKey("1", "1"), "1", "0", REASSIGNING)}),
                emptyArray());
    }

    @Test
    public void rebalanceShouldRemoveDeadSessions() {
        final Partition[] changeset = new ExactWeightRebalancer().apply(
                new Session[]{new Session("1", 1), new Session("2", 1)},
                new Partition[] {
                        new Partition(new PartitionKey("0", "0"), "0", null, ASSIGNED),
                        new Partition(new PartitionKey("0", "1"), "0", "1", REASSIGNING),
                        new Partition(new PartitionKey("1", "0"), "1", null, ASSIGNED),
                        new Partition(new PartitionKey("1", "1"), "0", null, ASSIGNED)});

        assertEquals(3, changeset.length);
        // All partitions must be in assigned state
        assertFalse(Stream.of(changeset).filter(p -> p.getState() != ASSIGNED).findAny().isPresent());
        // All partitions must not have nextSessionId
        assertFalse(Stream.of(changeset).filter(p -> p.getNextSession() != null).findAny().isPresent());

        assertFalse(Stream.of(changeset).filter(p->p.getKey().equals(new PartitionKey("1", "0"))).findAny().isPresent());
        assertEquals(1, Stream.of(changeset).filter(p -> p.getSession().equals("1")).count());
        assertEquals(2, Stream.of(changeset).filter(p -> p.getSession().equals("2")).count());
        assertFalse(Stream.of(changeset).filter(p -> p.getState() != ASSIGNED).findAny().isPresent());
    }

    @Test
    public void rebalanceShouldMoveToReassigningState() {
        final Partition[] changeset = new ExactWeightRebalancer().apply(
                new Session[]{new Session("1", 1), new Session("2", 1), new Session("3", 1)},
                new Partition[] {
                        new Partition(new PartitionKey("0", "0"), "1", null, ASSIGNED),
                        new Partition(new PartitionKey("0", "1"), "1", null, ASSIGNED),
                        new Partition(new PartitionKey("1", "0"), "2", null, ASSIGNED),
                        new Partition(new PartitionKey("1", "1"), "2", null, ASSIGNED)});
        assertEquals(1, changeset.length);
        final Partition changed = changeset[0];
        assertTrue(changed.getKey().equals(new PartitionKey("1", "0")) || changed.getKey().equals(new PartitionKey("1", "1")));
        assertEquals("2", changed.getSession());
        assertEquals("3", changed.getNextSession());
        Assert.assertEquals(REASSIGNING, changed.getState());
    }

    @Test
    public void rebalanceShouldTakeRebalancingPartitions() {
        final Partition[] changeset = new ExactWeightRebalancer().apply(
                new Session[]{new Session("1", 1), new Session("2", 1), new Session("3", 1)},
                new Partition[] {
                        new Partition(new PartitionKey("0", "0"), "1", null, ASSIGNED),
                        new Partition(new PartitionKey("0", "1"), "1", null, ASSIGNED),
                        new Partition(new PartitionKey("1", "0"), "2", null, REASSIGNING),
                        new Partition(new PartitionKey("1", "1"), "2", null, ASSIGNED)});
        assertEquals(1, changeset.length);
        final Partition changed = changeset[0];
        Assert.assertEquals(new PartitionKey("1", "0"), changed.getKey());
        assertEquals("2", changed.getSession());
        assertEquals("3", changed.getNextSession());
        Assert.assertEquals(REASSIGNING, changed.getState());
    }
}<commit_msg>ARUHA-305: fix static imports<commit_after>package org.zalando.nakadi.service.subscription;

import org.junit.Test;
import org.zalando.nakadi.service.subscription.model.Partition;
import org.zalando.nakadi.service.subscription.model.Partition.PartitionKey;
import org.zalando.nakadi.service.subscription.model.Session;

import java.util.stream.Stream;

import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.emptyArray;
import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
import static org.zalando.nakadi.service.subscription.model.Partition.State.ASSIGNED;
import static org.zalando.nakadi.service.subscription.model.Partition.State.REASSIGNING;

public class ExactWeightRebalancerTest {

    @Test(expected = IllegalArgumentException.class)
    public void splitByWeightShouldAcceptOnlyCorrectData1() {
        ExactWeightRebalancer.splitByWeight(1, new int[]{1, 1});
    }

    @Test(expected = IllegalArgumentException.class)
    public void splitByWeightShouldAcceptOnlyCorrectData2() {
        ExactWeightRebalancer.splitByWeight(2, new int[]{1, 0});
    }

    @Test
    public void splitByWeightMustCorrectlyWorkOnDifferentValues() {
        assertArrayEquals(
                new int[]{1, 1},
                ExactWeightRebalancer.splitByWeight(2, new int[]{1, 1}));
        assertArrayEquals(
                new int[]{2, 1},
                ExactWeightRebalancer.splitByWeight(3, new int[]{1, 1}));
        assertArrayEquals(
                new int[]{1, 2},
                ExactWeightRebalancer.splitByWeight(3, new int[]{1, 2}));
        assertArrayEquals(
                new int[]{34, 33, 33},
                ExactWeightRebalancer.splitByWeight(100, new int[]{1, 1, 1}));
        assertArrayEquals(
                new int[]{26, 25, 50},
                ExactWeightRebalancer.splitByWeight(101, new int[]{1, 1, 2}));
    }

    @Test
    public void rebalanceShouldHaveEmptyChangesetForBalancedData() {
        final ExactWeightRebalancer rebalancer = new ExactWeightRebalancer();

        // 1. Data contains only assigned
        final Session[] sessions = new Session[]{
                new Session("0", 1),
                new Session("1", 1)};
        assertThat(rebalancer.apply(sessions,
                new Partition[] {
                        new Partition(new PartitionKey("0", "0"), "0", null, ASSIGNED),
                        new Partition(new PartitionKey("0", "1"), "1", null, ASSIGNED),
                        new Partition(new PartitionKey("1", "0"), "1", null, ASSIGNED),
                        new Partition(new PartitionKey("1", "1"), "0", null, ASSIGNED)}),
                emptyArray());

        // 2. Data contains reassinging
        assertThat(rebalancer.apply(sessions,
                new Partition[] {
                        new Partition(new PartitionKey("0", "0"), "0", null, ASSIGNED),
                        new Partition(new PartitionKey("0", "1"), "0", "1", REASSIGNING),
                        new Partition(new PartitionKey("1", "0"), "1", null, ASSIGNED),
                        new Partition(new PartitionKey("1", "1"), "0", null, ASSIGNED)}),
                emptyArray());

        // 3. Data contains only reassinging
        assertThat(rebalancer.apply(sessions,
                new Partition[] {
                        new Partition(new PartitionKey("0", "0"), "0", "1", REASSIGNING),
                        new Partition(new PartitionKey("0", "1"), "0", "1", REASSIGNING),
                        new Partition(new PartitionKey("1", "0"), "1", "0", REASSIGNING),
                        new Partition(new PartitionKey("1", "1"), "1", "0", REASSIGNING)}),
                emptyArray());
    }

    @Test
    public void rebalanceShouldRemoveDeadSessions() {
        final Partition[] changeset = new ExactWeightRebalancer().apply(
                new Session[]{new Session("1", 1), new Session("2", 1)},
                new Partition[] {
                        new Partition(new PartitionKey("0", "0"), "0", null, ASSIGNED),
                        new Partition(new PartitionKey("0", "1"), "0", "1", REASSIGNING),
                        new Partition(new PartitionKey("1", "0"), "1", null, ASSIGNED),
                        new Partition(new PartitionKey("1", "1"), "0", null, ASSIGNED)});

        assertEquals(3, changeset.length);
        // All partitions must be in assigned state
        assertFalse(Stream.of(changeset).filter(p -> p.getState() != ASSIGNED).findAny().isPresent());
        // All partitions must not have nextSessionId
        assertFalse(Stream.of(changeset).filter(p -> p.getNextSession() != null).findAny().isPresent());

        assertFalse(Stream.of(changeset).filter(p->p.getKey().equals(new PartitionKey("1", "0"))).findAny().isPresent());
        assertEquals(1, Stream.of(changeset).filter(p -> p.getSession().equals("1")).count());
        assertEquals(2, Stream.of(changeset).filter(p -> p.getSession().equals("2")).count());
        assertFalse(Stream.of(changeset).filter(p -> p.getState() != ASSIGNED).findAny().isPresent());
    }

    @Test
    public void rebalanceShouldMoveToReassigningState() {
        final Partition[] changeset = new ExactWeightRebalancer().apply(
                new Session[]{new Session("1", 1), new Session("2", 1), new Session("3", 1)},
                new Partition[] {
                        new Partition(new PartitionKey("0", "0"), "1", null, ASSIGNED),
                        new Partition(new PartitionKey("0", "1"), "1", null, ASSIGNED),
                        new Partition(new PartitionKey("1", "0"), "2", null, ASSIGNED),
                        new Partition(new PartitionKey("1", "1"), "2", null, ASSIGNED)});
        assertEquals(1, changeset.length);
        final Partition changed = changeset[0];
        assertTrue(changed.getKey().equals(new PartitionKey("1", "0")) || changed.getKey().equals(new PartitionKey("1", "1")));
        assertEquals("2", changed.getSession());
        assertEquals("3", changed.getNextSession());
        assertEquals(REASSIGNING, changed.getState());
    }

    @Test
    public void rebalanceShouldTakeRebalancingPartitions() {
        final Partition[] changeset = new ExactWeightRebalancer().apply(
                new Session[]{new Session("1", 1), new Session("2", 1), new Session("3", 1)},
                new Partition[] {
                        new Partition(new PartitionKey("0", "0"), "1", null, ASSIGNED),
                        new Partition(new PartitionKey("0", "1"), "1", null, ASSIGNED),
                        new Partition(new PartitionKey("1", "0"), "2", null, REASSIGNING),
                        new Partition(new PartitionKey("1", "1"), "2", null, ASSIGNED)});
        assertEquals(1, changeset.length);
        final Partition changed = changeset[0];
        assertEquals(new PartitionKey("1", "0"), changed.getKey());
        assertEquals("2", changed.getSession());
        assertEquals("3", changed.getNextSession());
        assertEquals(REASSIGNING, changed.getState());
    }
}
{'fraction_non_alphanumeric': 0.14537873523280057, 'fraction_numerical': 0.023697011813759557, 'mean_word_length': 2.580741477979597, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6038479', 'n_tokens_mistral': 4206, 'n_tokens_neox': 3986, 'n_words': 896}