/*
 * Copyright 2015, Yahoo Inc.
 * Copyrights licensed under the New BSD License.
 * See the accompanying LICENSE file for terms.
 */

package com.yahoo.gondola.container;

import com.google.common.collect.Range;
import com.yahoo.gondola.Config;
import com.yahoo.gondola.Gondola;
import com.yahoo.gondola.Member;
import com.yahoo.gondola.Shard;
import com.yahoo.gondola.container.client.ShardManagerClient;

import org.mockito.Mock;
import org.mockito.MockitoAnnotations;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;

import java.util.Arrays;
import java.util.function.Consumer;

import static org.mockito.Matchers.any;
import static org.mockito.Matchers.anyInt;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.when;
import static org.testng.Assert.assertFalse;
import static org.testng.Assert.assertTrue;

public class ShardManagerTest {

    public static final String FROM_SHARD = "c1";
    public static final String TARGET_SHARD = "c2";
    public static final Integer FROM_MEMBER = 1;
    public static final int TIMEOUT_MS = 1000;
    ShardManager shardManager;

    @Mock
    RoutingFilter filter;

    @Mock
    LockManager lockManager;

    @Mock
    Config config;

    @Mock
    Config.ConfigMember configMember;

    @Mock
    ShardManagerClient shardManagerClient;

    @Mock
    Gondola gondola;

    @Mock
    Shard shard;

    @Mock
    Member member;

    @Mock
    Consumer<Member.SlaveStatus> updateCallback;

    @BeforeMethod
    public void setUp() throws Exception {
        MockitoAnnotations.initMocks(this);
        when(config.getMembersInShard(any())).thenReturn(Arrays.asList(configMember, configMember, configMember));
        when(configMember.getMemberId()).thenReturn(1,2,3);
        when(config.getMember(anyInt())).thenReturn(configMember);
        when(filter.isLeaderInShard(any())).thenReturn(true);
        when(filter.getGondola()).thenReturn(gondola);
        when(gondola.getShard(any())).thenReturn(shard);
        when(shard.getLocalMember()).thenReturn(member);
        shardManager = new ShardManager(filter, config, shardManagerClient);
    }

    @Test
    public void testStartObserving_success() throws Exception {
        callbackAnswer(true);
        assertFalse(shardManager.getObservedShards().contains(TARGET_SHARD));
        shardManager.startObserving(FROM_SHARD, TARGET_SHARD);
        assertTrue(shardManager.getObservedShards().contains(TARGET_SHARD));
        shardManager.startObserving(FROM_SHARD, TARGET_SHARD);
        assertTrue(shardManager.getObservedShards().contains(TARGET_SHARD));
    }

    @Test(expectedExceptions = ShardManagerProtocol.ShardManagerException.class)
    public void testStartObserving_failed() throws Exception {
        callbackAnswer(false);
        assertFalse(shardManager.getObservedShards().contains(TARGET_SHARD));
        shardManager.startObserving(FROM_SHARD, TARGET_SHARD);
    }


    @Test
    public void testStopObserving() throws Exception {
        callbackAnswer(true);

        shardManager.startObserving(FROM_SHARD, TARGET_SHARD);
        assertTrue(shardManager.getObservedShards().contains(TARGET_SHARD));
        shardManager.stopObserving(FROM_SHARD, TARGET_SHARD);
        assertFalse(shardManager.getObservedShards().contains(TARGET_SHARD));
        shardManager.stopObserving(FROM_SHARD, TARGET_SHARD);
        assertFalse(shardManager.getObservedShards().contains(TARGET_SHARD));
    }

    @Test
    public void testAssignBucket() throws Exception {
        Range<Integer> r = Range.closed(1, 2);
        shardManager.migrateBuckets(r, FROM_SHARD, TARGET_SHARD, TIMEOUT_MS);
    }

    @SuppressWarnings("unchecked")
    private void callbackAnswer(boolean running) {
        Member.SlaveStatus slaveStatus = new Member.SlaveStatus();
        slaveStatus.running = running;
        doAnswer(invocation -> {
            ((Consumer<Member.SlaveStatus>) invocation.getArguments()[1]).accept(slaveStatus);
            return null;
        }).when(member).setSlave(anyInt(), any());
    }
}
@@ -563,19 +563,36 @@
 ava.
-util.Arrays
+io.File;%0Aimport java.net.URL
 ;%0Aim
@@ -961,17 +961,21 @@
 HARD = %22
-c
+shard
 1%22;%0A    
@@ -1021,60 +1021,15 @@
  = %22
-c2%22;%0A    public static final Integer FROM_MEMBER = 1
+shard2%22
 ;%0A  
@@ -1190,31 +1190,154 @@
     
-@Mock%0A    Config config
+URL configUrl = ShardManagerTest.class.getClassLoader().getResource(%22gondola.conf%22);%0A    Config config = new Config(new File(configUrl.getFile()))
 ;%0A%0A 
@@ -1697,250 +1697,8 @@
 s);%0A
-        when(config.getMembersInShard(any())).thenReturn(Arrays.asList(configMember, configMember, configMember));%0A        when(configMember.getMemberId()).thenReturn(1,2,3);%0A        when(config.getMember(anyInt())).thenReturn(configMember);%0A
     
@@ -2105,20 +2105,30 @@
 kAnswer(
-true
+86, true, null
 );%0A     
@@ -2658,13 +2658,23 @@
 wer(
+86, 
 false
+, null
 );%0A 
@@ -2878,32 +2878,59 @@
 ows Exception %7B%0A
+        // Start observing%0A
         callback
@@ -2940,16 +2940,25 @@
 wer(
-true
+86, true, null
 );%0A
-%0A
     
@@ -3084,32 +3084,112 @@
 (TARGET_SHARD));
+%0A%0A        // Stop observing successfully%0A        callbackAnswer(-1, true, null);
 %0A        shardMa
@@ -3304,32 +3304,105 @@
 (TARGET_SHARD));
+%0A%0A        // Stop observing again%0A        callbackAnswer(-1, true, null);
 %0A        shardMa
@@ -3796,16 +3796,30 @@
 kAnswer(
+int memberId, 
 boolean 
@@ -3825,16 +3825,29 @@
  running
+, Exception e
 ) %7B%0A    
@@ -3869,16 +3869,120 @@
 eStatus 
+slaveStatus;%0A        if (memberId == -1) %7B%0A            slaveStatus = null;%0A        %7D else %7B%0A            
 slaveSta
@@ -4013,16 +4013,20 @@
 atus();%0A
+    
         
@@ -4056,16 +4056,110 @@
 unning;%0A
+            slaveStatus.memberId = memberId;%0A            slaveStatus.exception = e;%0A        %7D%0A
         
@@ -4358,11 +4358,75 @@
 ));%0A
+        when(member.getSlaveUpdate()).thenReturn(slaveStatus);%0A
     %7D%0A
-
 %7D
+%0A

Fix unit test for shardmanager