/*
 * Copyright 2018-present Open Networking Foundation
 *
 * 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 org.onosproject.segmentrouting;

import org.junit.Before;
import org.junit.Test;
import org.onlab.packet.IpAddress;
import org.onosproject.cluster.ClusterService;
import org.onosproject.cluster.DefaultControllerNode;
import org.onosproject.cluster.Leader;
import org.onosproject.cluster.Leadership;
import org.onosproject.cluster.LeadershipService;
import org.onosproject.cluster.NodeId;
import org.onosproject.net.DeviceId;
import org.onosproject.net.device.DeviceService;
import org.onosproject.store.service.StorageService;
import org.onosproject.store.service.TestConsistentMap;
import org.onosproject.store.service.TestConsistentMultimap;

import java.util.List;
import java.util.Optional;

import static org.easymock.EasyMock.createMock;
import static org.easymock.EasyMock.expect;
import static org.easymock.EasyMock.replay;
import static org.easymock.EasyMock.reset;
import static org.junit.Assert.*;
import static org.onosproject.segmentrouting.DefaultRoutingHandler.HASH_FUNCTION;

public class DefaultRoutingHandlerTest {
    private SegmentRoutingManager srManager;
    private DefaultRoutingHandler dfh;
    private MockWorkPartitionService mockWps;

    private static final DeviceId DEV1A = DeviceId.deviceId("of:1a");
    private static final DeviceId DEV1B = DeviceId.deviceId("of:1b");
    private static final DeviceId DEV2 = DeviceId.deviceId("of:2");

    private static final EdgePair EDGE_PAIR_1 = new EdgePair(DEV1A, DEV1B);
    private static final EdgePair EDGE_PAIR_2 = new EdgePair(DEV1B, DEV1A);
    private static final EdgePair EDGE_PAIR_3 = new EdgePair(DEV2, DeviceId.NONE);

    private static final NodeId NODE1 = NodeId.nodeId("192.168.1.1");
    private static final NodeId NODE2 = NodeId.nodeId("192.168.1.2");
    private static final NodeId NODE3 = NodeId.nodeId("192.168.1.3");
    private static final IpAddress IP1 = IpAddress.valueOf("192.168.1.1");
    private static final IpAddress IP2 = IpAddress.valueOf("192.168.1.2");
    private static final IpAddress IP3 = IpAddress.valueOf("192.168.1.3");

    /* This is generated by manually applying the hash function.
       It depends on the number of partitions defined in the MockWorkPartitionService */
    private static final String DEV1A_PARTITION_ID = String.valueOf(Math.abs(HASH_FUNCTION.apply(
            EDGE_PAIR_1).intValue()) % MockWorkPartitionService.NUM_PARTITIONS);
    private static final String DEV1B_PARTITION_ID = String.valueOf(Math.abs(HASH_FUNCTION.apply(
            EDGE_PAIR_2).intValue()) % MockWorkPartitionService.NUM_PARTITIONS);
    private static final String DEV2_PARTITION_ID = String.valueOf(Math.abs(HASH_FUNCTION.apply(
            EDGE_PAIR_3).intValue()) % MockWorkPartitionService.NUM_PARTITIONS);

    @Before
    public void setUp() {
        srManager = createMock(SegmentRoutingManager.class);
        srManager.storageService = createMock(StorageService.class);
        expect(srManager.storageService.consistentMapBuilder()).andReturn(
                new TestConsistentMap.Builder<>()).anyTimes();
        expect(srManager.storageService.consistentMultimapBuilder()).andReturn(
                new TestConsistentMultimap.Builder<>()).anyTimes();
        replay(srManager.storageService);
        srManager.routingRulePopulator = createMock(RoutingRulePopulator.class);
        srManager.deviceService = createMock(DeviceService.class);
        srManager.deviceConfiguration = createMock(DeviceConfiguration.class);
        mockWps = new MockWorkPartitionService();
        mockWps.leadershipService = createMock(LeadershipService.class);
        srManager.workPartitionService = mockWps;
        srManager.clusterService = createMock(ClusterService.class);
        dfh = new DefaultRoutingHandler(srManager);
    }

    private void clearShouldProgram() {
        dfh.invalidateShouldProgram(DEV1A);
        dfh.invalidateShouldProgram(DEV1B);
        dfh.invalidateShouldProgram(DEV2);
    }

    // Node 1 is the leader of switch 1A, 1B, and 2
    @Test
    public void testShouldHandleRoutingCase1() {
        expect(mockWps.leadershipService.getLeadership(DEV1A_PARTITION_ID)).andReturn(new Leadership(
                DEV1A_PARTITION_ID, new Leader(NODE1, 0, 0), List.of(NODE2, NODE3))).anyTimes();
        expect(mockWps.leadershipService.getLeadership(DEV1B_PARTITION_ID)).andReturn(new Leadership(
                DEV1B_PARTITION_ID, new Leader(NODE1, 0, 0), List.of(NODE2, NODE3))).anyTimes();
        expect(mockWps.leadershipService.getLeadership(DEV2_PARTITION_ID)).andReturn(new Leadership(
                DEV2_PARTITION_ID, new Leader(NODE1, 0, 0), List.of(NODE2, NODE3))).anyTimes();
        replay(mockWps.leadershipService);

        expect(srManager.getPairDeviceId(DEV1A)).andReturn(Optional.of(DEV1B)).anyTimes();
        expect(srManager.getPairDeviceId(DEV1B)).andReturn(Optional.of(DEV1A)).anyTimes();
        expect(srManager.getPairDeviceId(DEV2)).andReturn(Optional.empty()).anyTimes();
        replay(srManager);

        // Node 1 should program every device
        expect(srManager.clusterService.getLocalNode()).andReturn(new DefaultControllerNode(NODE1, IP1)).anyTimes();
        replay(srManager.clusterService);

        // shouldProgram is initially empty
        assertNull(dfh.shouldProgram.get(DEV1A));
        assertNull(dfh.shouldProgram.get(DEV1B));
        assertNull(dfh.shouldProgram.get(DEV2));
        assertTrue(dfh.shouldProgram(DEV1A));
        assertTrue(dfh.shouldProgram(DEV1B));
        assertTrue(dfh.shouldProgram(DEV2));
        assertEquals(NODE1, dfh.shouldProgram.get(DEV1A));
        assertEquals(NODE1, dfh.shouldProgram.get(DEV1B));
        assertEquals(NODE1, dfh.shouldProgram.get(DEV2));

        reset(srManager.clusterService);
        clearShouldProgram();

        // Node 2 should program no device
        expect(srManager.clusterService.getLocalNode()).andReturn(new DefaultControllerNode(NODE2, IP2)).anyTimes();
        replay(srManager.clusterService);
        assertNull(dfh.shouldProgram.get(DEV1A));
        assertNull(dfh.shouldProgram.get(DEV1B));
        assertNull(dfh.shouldProgram.get(DEV2));
        assertFalse(dfh.shouldProgram(DEV1A));
        assertFalse(dfh.shouldProgram(DEV1B));
        assertFalse(dfh.shouldProgram(DEV2));
        assertNotEquals(NODE2, dfh.shouldProgram.get(DEV1A));
        assertNotEquals(NODE2, dfh.shouldProgram.get(DEV1B));
        assertNotEquals(NODE2, dfh.shouldProgram.get(DEV2));

        reset(srManager.clusterService);
        clearShouldProgram();

        // Node 3 should program no device
        expect(srManager.clusterService.getLocalNode()).andReturn(new DefaultControllerNode(NODE3, IP3)).anyTimes();
        replay(srManager.clusterService);
        assertNull(dfh.shouldProgram.get(DEV1A));
        assertNull(dfh.shouldProgram.get(DEV1B));
        assertNull(dfh.shouldProgram.get(DEV2));
        assertFalse(dfh.shouldProgram(DEV1A));
        assertFalse(dfh.shouldProgram(DEV1B));
        assertFalse(dfh.shouldProgram(DEV2));
        assertNotEquals(NODE3, dfh.shouldProgram.get(DEV1A));
        assertNotEquals(NODE3, dfh.shouldProgram.get(DEV1B));
        assertNotEquals(NODE3, dfh.shouldProgram.get(DEV2));
    }

    /*
     * Node 1 is the leader of switch 1A, 1B
     * Node 2 is the leader of switch 2
     */
    @Test
    public void testShouldHandleRoutingCase2() {
        expect(mockWps.leadershipService.getLeadership(DEV1A_PARTITION_ID)).andReturn(new Leadership(
                DEV1A_PARTITION_ID, new Leader(NODE1, 0, 0), List.of(NODE2, NODE3))).anyTimes();
        expect(mockWps.leadershipService.getLeadership(DEV1B_PARTITION_ID)).andReturn(new Leadership(
                DEV1B_PARTITION_ID, new Leader(NODE1, 0, 0), List.of(NODE2, NODE3))).anyTimes();
        expect(mockWps.leadershipService.getLeadership(DEV2_PARTITION_ID)).andReturn(new Leadership(
                DEV2_PARTITION_ID, new Leader(NODE2, 0, 0), List.of(NODE2, NODE3))).anyTimes();
        replay(mockWps.leadershipService);

        expect(srManager.getPairDeviceId(DEV1A)).andReturn(Optional.of(DEV1B)).anyTimes();
        expect(srManager.getPairDeviceId(DEV1B)).andReturn(Optional.of(DEV1A)).anyTimes();
        expect(srManager.getPairDeviceId(DEV2)).andReturn(Optional.empty()).anyTimes();
        replay(srManager);

        // Node 1 should program 1A, 1B
        expect(srManager.clusterService.getLocalNode()).andReturn(new DefaultControllerNode(NODE1, IP1)).anyTimes();
        replay(srManager.clusterService);
        assertNull(dfh.shouldProgram.get(DEV1A));
        assertNull(dfh.shouldProgram.get(DEV1B));
        assertNull(dfh.shouldProgram.get(DEV2));
        assertTrue(dfh.shouldProgram(DEV1A));
        assertTrue(dfh.shouldProgram(DEV1B));
        assertFalse(dfh.shouldProgram(DEV2));
        assertEquals(NODE1, dfh.shouldProgram.get(DEV1A));
        assertEquals(NODE1, dfh.shouldProgram.get(DEV1B));
        assertNotEquals(NODE1, dfh.shouldProgram.get(DEV2));

        reset(srManager.clusterService);
        clearShouldProgram();

        // Node 2 should program 2
        expect(srManager.clusterService.getLocalNode()).andReturn(new DefaultControllerNode(NODE2, IP2)).anyTimes();
        replay(srManager.clusterService);
        assertNull(dfh.shouldProgram.get(DEV1A));
        assertNull(dfh.shouldProgram.get(DEV1B));
        assertNull(dfh.shouldProgram.get(DEV2));
        assertFalse(dfh.shouldProgram(DEV1A));
        assertFalse(dfh.shouldProgram(DEV1B));
        assertTrue(dfh.shouldProgram(DEV2));
        assertNotEquals(NODE2, dfh.shouldProgram.get(DEV1A));
        assertNotEquals(NODE2, dfh.shouldProgram.get(DEV1B));
        assertEquals(NODE2, dfh.shouldProgram.get(DEV2));

        reset(srManager.clusterService);
        clearShouldProgram();

        // Node 3 should program no device
        expect(srManager.clusterService.getLocalNode()).andReturn(new DefaultControllerNode(NODE3, IP3)).anyTimes();
        replay(srManager.clusterService);
        assertNull(dfh.shouldProgram.get(DEV1A));
        assertNull(dfh.shouldProgram.get(DEV1B));
        assertNull(dfh.shouldProgram.get(DEV2));
        assertFalse(dfh.shouldProgram(DEV1A));
        assertFalse(dfh.shouldProgram(DEV1B));
        assertFalse(dfh.shouldProgram(DEV2));
        assertNotEquals(NODE3, dfh.shouldProgram.get(DEV1A));
        assertNotEquals(NODE3, dfh.shouldProgram.get(DEV1B));
        assertNotEquals(NODE3, dfh.shouldProgram.get(DEV2));
    }

    /*
     * Node 1 is the leader of switch 1A, 1B
     * Node 3 is the leader of switch 2
     */
    @Test
    public void testShouldHandleRoutingCase3() {
        expect(mockWps.leadershipService.getLeadership(DEV1A_PARTITION_ID)).andReturn(new Leadership(
                DEV1A_PARTITION_ID, new Leader(NODE1, 0, 0), List.of(NODE2, NODE3))).anyTimes();
        expect(mockWps.leadershipService.getLeadership(DEV1B_PARTITION_ID)).andReturn(new Leadership(
                DEV1B_PARTITION_ID, new Leader(NODE1, 0, 0), List.of(NODE2, NODE3))).anyTimes();
        expect(mockWps.leadershipService.getLeadership(DEV2_PARTITION_ID)).andReturn(new Leadership(
                DEV2_PARTITION_ID, new Leader(NODE3, 0, 0), List.of(NODE2, NODE3))).anyTimes();
        replay(mockWps.leadershipService);

        expect(srManager.getPairDeviceId(DEV1A)).andReturn(Optional.of(DEV1B)).anyTimes();
        expect(srManager.getPairDeviceId(DEV1B)).andReturn(Optional.of(DEV1A)).anyTimes();
        expect(srManager.getPairDeviceId(DEV2)).andReturn(Optional.empty()).anyTimes();
        replay(srManager);

        // Node 1 should program 1A, 1B
        expect(srManager.clusterService.getLocalNode()).andReturn(new DefaultControllerNode(NODE1, IP1)).anyTimes();
        replay(srManager.clusterService);
        assertNull(dfh.shouldProgram.get(DEV1A));
        assertNull(dfh.shouldProgram.get(DEV1B));
        assertNull(dfh.shouldProgram.get(DEV2));
        assertTrue(dfh.shouldProgram(DEV1A));
        assertTrue(dfh.shouldProgram(DEV1B));
        assertFalse(dfh.shouldProgram(DEV2));
        assertEquals(NODE1, dfh.shouldProgram.get(DEV1A));
        assertEquals(NODE1, dfh.shouldProgram.get(DEV1B));
        assertNotEquals(NODE1, dfh.shouldProgram.get(DEV2));

        reset(srManager.clusterService);
        clearShouldProgram();

        // Node 2 should program no device
        expect(srManager.clusterService.getLocalNode()).andReturn(new DefaultControllerNode(NODE2, IP2)).anyTimes();
        replay(srManager.clusterService);
        assertNull(dfh.shouldProgram.get(DEV1A));
        assertNull(dfh.shouldProgram.get(DEV1B));
        assertNull(dfh.shouldProgram.get(DEV2));
        assertFalse(dfh.shouldProgram(DEV1A));
        assertFalse(dfh.shouldProgram(DEV1B));
        assertFalse(dfh.shouldProgram(DEV2));
        assertNotEquals(NODE2, dfh.shouldProgram.get(DEV1A));
        assertNotEquals(NODE2, dfh.shouldProgram.get(DEV1B));
        assertNotEquals(NODE2, dfh.shouldProgram.get(DEV2));

        reset(srManager.clusterService);
        clearShouldProgram();

        // Node 3 should program 2
        expect(srManager.clusterService.getLocalNode()).andReturn(new DefaultControllerNode(NODE3, IP3)).anyTimes();
        replay(srManager.clusterService);
        assertNull(dfh.shouldProgram.get(DEV1A));
        assertNull(dfh.shouldProgram.get(DEV1B));
        assertNull(dfh.shouldProgram.get(DEV2));
        assertFalse(dfh.shouldProgram(DEV1A));
        assertFalse(dfh.shouldProgram(DEV1B));
        assertTrue(dfh.shouldProgram(DEV2));
        assertNotEquals(NODE3, dfh.shouldProgram.get(DEV1A));
        assertNotEquals(NODE3, dfh.shouldProgram.get(DEV1B));
        assertEquals(NODE3, dfh.shouldProgram.get(DEV2));
    }

    /*
     * Node 2 is the leader of switch 1A, 1B and Node 3 is the leader of 2.
     * Later on, node 1 becomes the leader of 1A, 1B
     */
    @Test
    public void testShouldHandleRoutingCase4() {
        expect(mockWps.leadershipService.getLeadership(DEV1A_PARTITION_ID)).andReturn(new Leadership(
                DEV1A_PARTITION_ID, new Leader(NODE2, 0, 0), List.of(NODE2, NODE3))).anyTimes();
        expect(mockWps.leadershipService.getLeadership(DEV1B_PARTITION_ID)).andReturn(new Leadership(
                DEV1B_PARTITION_ID, new Leader(NODE2, 0, 0), List.of(NODE2, NODE3))).anyTimes();
        expect(mockWps.leadershipService.getLeadership(DEV2_PARTITION_ID)).andReturn(new Leadership(
                DEV2_PARTITION_ID, new Leader(NODE3, 0, 0), List.of(NODE2, NODE3))).anyTimes();
        replay(mockWps.leadershipService);

        expect(srManager.getPairDeviceId(DEV1A)).andReturn(Optional.of(DEV1B)).anyTimes();
        expect(srManager.getPairDeviceId(DEV1B)).andReturn(Optional.of(DEV1A)).anyTimes();
        expect(srManager.getPairDeviceId(DEV2)).andReturn(Optional.empty()).anyTimes();
        replay(srManager);

        // Node 1 should program no device
        expect(srManager.clusterService.getLocalNode()).andReturn(new DefaultControllerNode(NODE1, IP1)).anyTimes();
        replay(srManager.clusterService);
        assertNull(dfh.shouldProgram.get(DEV1A));
        assertNull(dfh.shouldProgram.get(DEV1B));
        assertNull(dfh.shouldProgram.get(DEV2));
        assertFalse(dfh.shouldProgram(DEV1A));
        assertFalse(dfh.shouldProgram(DEV1B));
        assertFalse(dfh.shouldProgram(DEV2));
        assertNotEquals(NODE1, dfh.shouldProgram.get(DEV1A));
        assertNotEquals(NODE1, dfh.shouldProgram.get(DEV1B));
        assertNotEquals(NODE1, dfh.shouldProgram.get(DEV2));

        reset(srManager.clusterService);
        clearShouldProgram();

        // Node 2 should program 1A and 1B
        expect(srManager.clusterService.getLocalNode()).andReturn(new DefaultControllerNode(NODE2, IP2)).anyTimes();
        replay(srManager.clusterService);
        assertNull(dfh.shouldProgram.get(DEV1A));
        assertNull(dfh.shouldProgram.get(DEV1B));
        assertNull(dfh.shouldProgram.get(DEV2));
        assertTrue(dfh.shouldProgram(DEV1A));
        assertTrue(dfh.shouldProgram(DEV1B));
        assertFalse(dfh.shouldProgram(DEV2));
        assertEquals(NODE2, dfh.shouldProgram.get(DEV1A));
        assertEquals(NODE2, dfh.shouldProgram.get(DEV1B));
        assertNotEquals(NODE2, dfh.shouldProgram.get(DEV2));

        reset(srManager.clusterService);
        clearShouldProgram();

        // Node 3 should program 2
        expect(srManager.clusterService.getLocalNode()).andReturn(new DefaultControllerNode(NODE3, IP3)).anyTimes();
        replay(srManager.clusterService);
        assertNull(dfh.shouldProgram.get(DEV1A));
        assertNull(dfh.shouldProgram.get(DEV1B));
        assertNull(dfh.shouldProgram.get(DEV2));
        assertFalse(dfh.shouldProgram(DEV1A));
        assertFalse(dfh.shouldProgram(DEV1B));
        assertTrue(dfh.shouldProgram(DEV2));
        assertNotEquals(NODE3, dfh.shouldProgram.get(DEV1A));
        assertNotEquals(NODE3, dfh.shouldProgram.get(DEV1B));
        assertEquals(NODE3, dfh.shouldProgram.get(DEV2));

        // Partition of switch 1A moves to Node 1. This can happen for a cluster event
        reset(mockWps.leadershipService);
        expect(mockWps.leadershipService.getLeadership(DEV1A_PARTITION_ID)).andReturn(new Leadership(
                DEV1A_PARTITION_ID, new Leader(NODE1, 0, 0), List.of(NODE2, NODE3))).anyTimes();
        expect(mockWps.leadershipService.getLeadership(DEV1B_PARTITION_ID)).andReturn(new Leadership(
                DEV1B_PARTITION_ID, new Leader(NODE1, 0, 0), List.of(NODE2, NODE3))).anyTimes();
        expect(mockWps.leadershipService.getLeadership(DEV2_PARTITION_ID)).andReturn(new Leadership(
                DEV2_PARTITION_ID, new Leader(NODE3, 0, 0), List.of(NODE2, NODE3))).anyTimes();
        replay(mockWps.leadershipService);

        reset(srManager.clusterService);
        clearShouldProgram();

        // Node 1 should program 1A, 1B
        expect(srManager.clusterService.getLocalNode()).andReturn(new DefaultControllerNode(NODE1, IP1)).anyTimes();
        replay(srManager.clusterService);
        assertNull(dfh.shouldProgram.get(DEV1A));
        assertNull(dfh.shouldProgram.get(DEV1B));
        assertNull(dfh.shouldProgram.get(DEV2));
        assertTrue(dfh.shouldProgram(DEV1A));
        assertTrue(dfh.shouldProgram(DEV1B));
        assertFalse(dfh.shouldProgram(DEV2));
        assertEquals(NODE1, dfh.shouldProgram.get(DEV1A));
        assertEquals(NODE1, dfh.shouldProgram.get(DEV1B));
        assertNotEquals(NODE1, dfh.shouldProgram.get(DEV2));

        reset(srManager.clusterService);
        clearShouldProgram();

        // Node 2 should program no device
        expect(srManager.clusterService.getLocalNode()).andReturn(new DefaultControllerNode(NODE2, IP2)).anyTimes();
        replay(srManager.clusterService);
        assertNull(dfh.shouldProgram.get(DEV1A));
        assertNull(dfh.shouldProgram.get(DEV1B));
        assertNull(dfh.shouldProgram.get(DEV2));
        assertFalse(dfh.shouldProgram(DEV1A));
        assertFalse(dfh.shouldProgram(DEV1B));
        assertFalse(dfh.shouldProgram(DEV2));
        assertNotEquals(NODE2, dfh.shouldProgram.get(DEV1A));
        assertNotEquals(NODE2, dfh.shouldProgram.get(DEV1B));
        assertNotEquals(NODE2, dfh.shouldProgram.get(DEV2));

        reset(srManager.clusterService);
        clearShouldProgram();

        // Node 3 should program 2
        expect(srManager.clusterService.getLocalNode()).andReturn(new DefaultControllerNode(NODE3, IP3)).anyTimes();
        replay(srManager.clusterService);
        assertNull(dfh.shouldProgram.get(DEV1A));
        assertNull(dfh.shouldProgram.get(DEV1B));
        assertNull(dfh.shouldProgram.get(DEV2));
        assertFalse(dfh.shouldProgram(DEV1A));
        assertFalse(dfh.shouldProgram(DEV1B));
        assertTrue(dfh.shouldProgram(DEV2));
        assertNotEquals(NODE3, dfh.shouldProgram.get(DEV1A));
        assertNotEquals(NODE3, dfh.shouldProgram.get(DEV1B));
        assertEquals(NODE3, dfh.shouldProgram.get(DEV2));
    }

    /*
     * Node 1 is the leader of 1A, 1B. Node 2 becomes the leader of 1A, 1B later
     * shouldP is not purged in time. This can easily happen if we dont get in
     * time cluster/mastership events. shouldProgram absorbs this negative scenario.
     */
    @Test
    public void testShouldHandleRoutingCase5() {
        expect(mockWps.leadershipService.getLeadership(DEV1A_PARTITION_ID)).andReturn(new Leadership(
                DEV1A_PARTITION_ID, new Leader(NODE1, 0, 0), List.of(NODE2, NODE3))).anyTimes();
        expect(mockWps.leadershipService.getLeadership(DEV1B_PARTITION_ID)).andReturn(new Leadership(
                DEV1B_PARTITION_ID, new Leader(NODE1, 0, 0), List.of(NODE2, NODE3))).anyTimes();
        replay(mockWps.leadershipService);

        expect(srManager.getPairDeviceId(DEV1A)).andReturn(Optional.of(DEV1B)).anyTimes();
        expect(srManager.getPairDeviceId(DEV1B)).andReturn(Optional.of(DEV1A)).anyTimes();
        replay(srManager);

        // Node 1 should program both 1A and 1B
        expect(srManager.clusterService.getLocalNode()).andReturn(new DefaultControllerNode(NODE1, IP1)).anyTimes();
        replay(srManager.clusterService);
        assertNull(dfh.shouldProgram.get(DEV1A));
        assertNull(dfh.shouldProgram.get(DEV1B));
        assertTrue(dfh.shouldProgram(DEV1A));
        assertTrue(dfh.shouldProgram(DEV1B));
        assertEquals(NODE1, dfh.shouldProgram.get(DEV1A));
        assertEquals(NODE1, dfh.shouldProgram.get(DEV1B));

        reset(srManager.clusterService);
        clearShouldProgram();

        // Node 2 should program no device
        expect(srManager.clusterService.getLocalNode()).andReturn(new DefaultControllerNode(NODE2, IP2)).anyTimes();
        replay(srManager.clusterService);
        assertNull(dfh.shouldProgram.get(DEV1A));
        assertNull(dfh.shouldProgram.get(DEV1B));
        assertFalse(dfh.shouldProgram(DEV1A));
        assertFalse(dfh.shouldProgram(DEV1B));
        assertNotEquals(NODE2, dfh.shouldProgram.get(DEV1A));
        assertNotEquals(NODE2, dfh.shouldProgram.get(DEV1B));

        // Leadership moves to Node 2
        reset(mockWps.leadershipService);
        expect(mockWps.leadershipService.getLeadership(DEV1A_PARTITION_ID)).andReturn(new Leadership(
                DEV1A_PARTITION_ID, new Leader(NODE2, 0, 0), List.of(NODE2, NODE3))).anyTimes();
        expect(mockWps.leadershipService.getLeadership(DEV1B_PARTITION_ID)).andReturn(new Leadership(
                DEV1B_PARTITION_ID, new Leader(NODE2, 0, 0), List.of(NODE2, NODE3))).anyTimes();
        replay(mockWps.leadershipService);

        reset(srManager.clusterService);

        // Node 1 should program 1A, 1B
        expect(srManager.clusterService.getLocalNode()).andReturn(new DefaultControllerNode(NODE1, IP1)).anyTimes();
        replay(srManager.clusterService);
        assertNotNull(dfh.shouldProgram.get(DEV1A));
        assertNotNull(dfh.shouldProgram.get(DEV1B));
        assertTrue(dfh.shouldProgram(DEV1A));
        assertTrue(dfh.shouldProgram(DEV1B));
        assertEquals(NODE1, dfh.shouldProgram.get(DEV1A));
        assertEquals(NODE1, dfh.shouldProgram.get(DEV1B));

        reset(srManager.clusterService);

        // Node 2 should program no device
        expect(srManager.clusterService.getLocalNode()).andReturn(new DefaultControllerNode(NODE2, IP2)).anyTimes();
        replay(srManager.clusterService);
        assertNotNull(dfh.shouldProgram.get(DEV1A));
        assertNotNull(dfh.shouldProgram.get(DEV1B));
        assertFalse(dfh.shouldProgram(DEV1A));
        assertFalse(dfh.shouldProgram(DEV1B));
        assertEquals(NODE1, dfh.shouldProgram.get(DEV1A));
        assertEquals(NODE1, dfh.shouldProgram.get(DEV1B));
    }

    // There is no leadership for 1A, 1B. Super-damaged cluster.
    @Test(expected = NullPointerException.class)
    public void testShouldHandleRoutingCase6() {
        expect(mockWps.leadershipService.getLeadership(DEV1A_PARTITION_ID)).andReturn(null).anyTimes();
        expect(mockWps.leadershipService.getLeadership(DEV1B_PARTITION_ID)).andReturn(null).anyTimes();
        replay(mockWps.leadershipService);

        expect(srManager.getPairDeviceId(DEV1A)).andReturn(Optional.of(DEV1B)).anyTimes();
        expect(srManager.getPairDeviceId(DEV1B)).andReturn(Optional.of(DEV1A)).anyTimes();
        replay(srManager);

        // Node 1 should program no device
        expect(srManager.clusterService.getLocalNode()).andReturn(new DefaultControllerNode(NODE1, IP1)).anyTimes();
        replay(srManager.clusterService);
        assertNull(dfh.shouldProgram.get(DEV1A));
        assertNull(dfh.shouldProgram.get(DEV1B));
        assertFalse(dfh.shouldProgram(DEV1A));
        assertFalse(dfh.shouldProgram(DEV1B));
        assertNull(dfh.shouldProgram.get(DEV1A));
        assertNull(dfh.shouldProgram.get(DEV1B));

        reset(srManager.clusterService);
        clearShouldProgram();

        // Node 2 should program no device
        expect(srManager.clusterService.getLocalNode()).andReturn(new DefaultControllerNode(NODE2, IP2)).anyTimes();
        replay(srManager.clusterService);
        assertNull(dfh.shouldProgram.get(DEV1A));
        assertNull(dfh.shouldProgram.get(DEV1B));
        assertFalse(dfh.shouldProgram(DEV1A));
        assertFalse(dfh.shouldProgram(DEV1B));
        assertNull(dfh.shouldProgram.get(DEV1A));
        assertNull(dfh.shouldProgram.get(DEV1B));

        assertFalse(dfh.shouldProgram.containsKey(DEV1A));
        assertFalse(dfh.shouldProgram.containsKey(DEV1B));
    }
}