/*
 * Copyright 2016-2018 Seznam.cz, a.s.
 *
 * 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 cz.seznam.euphoria.kafka;

import org.apache.kafka.clients.consumer.Consumer;
import org.junit.Test;

import java.util.Collections;

import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;


public class KafkaSourceTest {

  /**
   * Test if where there are no partitions
   * an IllegalStateException is expected
   */
  @Test(expected = IllegalStateException.class)
  @SuppressWarnings("unchecked")
  public void testPartitions() {
    KafkaSource source = mock(KafkaSource.class);
    Consumer<byte[], byte[]> consumer = mock(Consumer.class);
    when(consumer.partitionsFor(any(String.class))).thenReturn(Collections.emptyList());
    when(source.newConsumer(any(), any(), any())).thenReturn(consumer);
    when(source.getPartitions()).thenCallRealMethod();
    source.getPartitions();
  }
}
@@ -698,18 +698,154 @@
 org.
-junit.Test
+apache.kafka.common.Node;%0Aimport org.apache.kafka.common.PartitionInfo;%0Aimport org.junit.Test;%0Aimport org.junit.experimental.categories.Categories
 ;%0A%0Ai
@@ -872,16 +872,39 @@
 ections;
+%0Aimport java.util.List;
 %0A%0Aimport
@@ -1255,16 +1255,18 @@
 oid test
+No
 Partitio
@@ -1272,16 +1272,732 @@
 ons() %7B%0A
+    tryGetPartitions(Collections.emptyList());%0A  %7D%0A%0A  @SuppressWarnings(%22unchecked%22)%0A  public void testPartitions() %7B%0A    Node leaderNode = new Node(1, %22localhost%22, 3333);%0A    PartitionInfo pi = new PartitionInfo(%22topic%22, 0, leaderNode, null, null);%0A    tryGetPartitions(Collections.emptyList());%0A  %7D%0A%0A  @Test(expected = IllegalStateException.class)%0A  @SuppressWarnings(%22unchecked%22)%0A  public void testNoLeader() %7B%0A    Node leaderNode = new Node(-1, %22localhost%22, 3333); // -1 no leader, constant not even defined in Kafka%0A    PartitionInfo pi = new PartitionInfo(%22topic%22, 0, leaderNode, null, null);%0A    tryGetPartitions(Collections.emptyList());%0A  %7D%0A%0A  public void tryGetPartitions(List%3CPartitionInfo%3E partitions) %7B%0A
     Kafk
@@ -2163,39 +2163,26 @@
 nReturn(
-Collections.emptyList()
+partitions
 );%0A    w

Add tests in KafkaSourceTest