// 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 com.cloud.network;

import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyBoolean;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.lenient;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Date;
import java.util.List;
import java.util.Vector;

import com.cloud.network.dao.PublicIpQuarantineDao;
import com.cloud.network.vo.PublicIpQuarantineVO;
import com.cloud.user.Account;
import com.cloud.user.AccountManager;
import org.apache.cloudstack.context.CallContext;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.InjectMocks;
import org.mockito.Mock;
import org.mockito.MockedStatic;
import org.mockito.Mockito;
import org.mockito.Spy;
import org.mockito.junit.MockitoJUnitRunner;

import com.cloud.exception.ResourceUnavailableException;
import com.cloud.network.Network.Service;
import com.cloud.network.dao.IPAddressDao;
import com.cloud.network.dao.IPAddressVO;
import com.cloud.network.dao.NetworkDao;
import com.cloud.network.dao.NetworkVO;
import com.cloud.network.rules.FirewallRule;
import com.cloud.network.rules.StaticNat;
import com.cloud.network.rules.StaticNatImpl;
import com.cloud.network.vpc.VpcManager;
import com.cloud.dc.dao.VlanDao;
import com.cloud.dc.VlanVO;
import com.cloud.offerings.NetworkOfferingVO;
import com.cloud.offerings.dao.NetworkOfferingDao;
import com.cloud.user.AccountVO;
import com.cloud.utils.net.Ip;

@RunWith(MockitoJUnitRunner.class)
public class IpAddressManagerTest {

    @Mock
    IPAddressDao ipAddressDao;

    @Mock
    NetworkDao networkDao;

    @Mock
    NetworkOfferingDao networkOfferingDao;

    @Mock
    IPAddressVO ipAddressVoMock;

    @Spy
    @InjectMocks
    IpAddressManagerImpl ipAddressManager;

    @InjectMocks
    NetworkModelImpl networkModel = Mockito.spy(new NetworkModelImpl());

    IPAddressVO ipAddressVO;

    AccountVO account;

    @Mock
    PublicIpQuarantineVO publicIpQuarantineVOMock;

    @Mock
    PublicIpQuarantineDao publicIpQuarantineDaoMock;

    @Mock
    IpAddress ipAddressMock;

    @Mock
    AccountVO newOwnerMock;

    @Mock
    AccountVO previousOwnerMock;

    @Mock
    AccountManager accountManagerMock;

    @Mock
    VpcManager vpcMgr;

    @Mock
    VlanDao vlanDao;

    final long dummyID = 1L;

    final String UUID = "uuid";

    private static final Date currentDate = new Date(100L);

    private static final Date beforeCurrentDate = new Date(99L);

    private static final Date afterCurrentDate = new Date(101L);

    @Before
    public void setup() throws ResourceUnavailableException {

        ipAddressVO = new IPAddressVO(new Ip("192.0.0.1"), 1L, 1L, 1L,false);
        ipAddressVO.setAllocatedToAccountId(1L);

        account = new AccountVO("admin", 1L, null, Account.Type.ADMIN, 1L, "c65a73d5-ebbd-11e7-8f45-107b44277808");
        account.setId(1L);

        NetworkOfferingVO networkOfferingVO = Mockito.mock(NetworkOfferingVO.class);
        networkOfferingVO.setSharedSourceNat(false);

        Mockito.when(networkOfferingDao.findById(Mockito.anyLong())).thenReturn(networkOfferingVO);
    }

    @Test
    public void testGetStaticNatSourceIps() {
        String publicIpAddress = "192.168.1.3";
        IPAddressVO vo = mock(IPAddressVO.class);
        lenient().when(vo.getAddress()).thenReturn(new Ip(publicIpAddress));
        lenient().when(vo.getId()).thenReturn(1l);

        when(ipAddressDao.findById(anyLong())).thenReturn(vo);
        StaticNat snat = new StaticNatImpl(1, 1, 1, 1, publicIpAddress, false);

        List<IPAddressVO> ips = ipAddressManager.getStaticNatSourceIps(Collections.singletonList(snat));
        Assert.assertNotNull(ips);
        Assert.assertEquals(1, ips.size());

        IPAddressVO returnedVO = ips.get(0);
        Assert.assertEquals(vo, returnedVO);
    }

    @Test
    public void isIpEqualsGatewayOrNetworkOfferingsEmptyTestRequestedIpEqualsIp6Gateway() {
        Network network = setTestIsIpEqualsGatewayOrNetworkOfferingsEmpty(0l, "gateway", "ip6Gateway", null, new ArrayList<Service>());

        boolean result = ipAddressManager.isIpEqualsGatewayOrNetworkOfferingsEmpty(network, "ip6Gateway");

        Mockito.verify(networkModel, Mockito.times(0)).listNetworkOfferingServices(Mockito.anyLong());
        Assert.assertTrue(result);
    }

    @Test
    public void isIpEqualsGatewayOrNetworkOfferingsEmptyTestRequestedIpEqualsGateway() {
        Network network = setTestIsIpEqualsGatewayOrNetworkOfferingsEmpty(0l, "gateway", "ip6Gateway", null, new ArrayList<Service>());

        boolean result = ipAddressManager.isIpEqualsGatewayOrNetworkOfferingsEmpty(network, "gateway");

        Mockito.verify(networkModel, Mockito.times(0)).listNetworkOfferingServices(Mockito.anyLong());
        Assert.assertTrue(result);
    }

    @Test
    public void isIpEqualsGatewayOrNetworkOfferingsEmptyTestExpectFalseServicesNotEmpty() {
        List<Service> services = new ArrayList<Service>();
        Service serviceGateway = new Service("Gateway");
        services.add(serviceGateway);
        Network network = setTestIsIpEqualsGatewayOrNetworkOfferingsEmpty(0l, "gateway", "ip6Gateway", null, services);

        boolean result = ipAddressManager.isIpEqualsGatewayOrNetworkOfferingsEmpty(network, "requestedIp");

        Mockito.verify(networkModel).listNetworkOfferingServices(Mockito.anyLong());
        Assert.assertFalse(result);
    }

    @Test
    public void isIpEqualsGatewayOrNetworkOfferingsEmptyTestExpectFalseServicesCidrNotNull() {
        Network network = setTestIsIpEqualsGatewayOrNetworkOfferingsEmpty(0l, "gateway", "ip6Gateway", "cidr", new ArrayList<Service>());

        boolean result = ipAddressManager.isIpEqualsGatewayOrNetworkOfferingsEmpty(network, "requestedIp");

        Mockito.verify(networkModel).listNetworkOfferingServices(Mockito.anyLong());
        Assert.assertFalse(result);
    }

    @Test
    public void assertSourceNatImplementedNetwork() {

        NetworkVO networkImplemented = Mockito.mock(NetworkVO.class);
        lenient().when(networkImplemented.getTrafficType()).thenReturn(Networks.TrafficType.Guest);
        lenient().when(networkImplemented.getNetworkOfferingId()).thenReturn(8L);
        lenient().when(networkImplemented.getState()).thenReturn(Network.State.Implemented);
        when(networkImplemented.getGuestType()).thenReturn(Network.GuestType.Isolated);
        when(networkImplemented.getVpcId()).thenReturn(null);
        when(networkImplemented.getId()).thenReturn(1L);

        Mockito.lenient().when(networkDao.findById(1L)).thenReturn(networkImplemented);
        doReturn(null).when(ipAddressManager).getExistingSourceNatInNetwork(1L, 1L);

        boolean isSourceNat = ipAddressManager.isSourceNatAvailableForNetwork(account, ipAddressVO, networkImplemented);

        assertTrue("Source NAT should be true", isSourceNat);
    }

    @Test
    public void assertSourceNatAllocatedNetwork() {

        NetworkVO networkAllocated = Mockito.mock(NetworkVO.class);
        lenient().when(networkAllocated.getTrafficType()).thenReturn(Networks.TrafficType.Guest);
        when(networkAllocated.getNetworkOfferingId()).thenReturn(8L);
        lenient().when(networkAllocated.getState()).thenReturn(Network.State.Allocated);
        when(networkAllocated.getGuestType()).thenReturn(Network.GuestType.Isolated);
        when(networkAllocated.getVpcId()).thenReturn(null);
        when(networkAllocated.getId()).thenReturn(2L);

        Mockito.lenient().when(networkDao.findById(2L)).thenReturn(networkAllocated);
        doReturn(null).when(ipAddressManager).getExistingSourceNatInNetwork(1L, 2L);

        assertTrue(ipAddressManager.isSourceNatAvailableForNetwork(account, ipAddressVO, networkAllocated));
    }

    @Test
    public void assertExistingSourceNatAllocatedNetwork() {

        NetworkVO networkNat = Mockito.mock(NetworkVO.class);
        lenient().when(networkNat.getTrafficType()).thenReturn(Networks.TrafficType.Guest);
        when(networkNat.getNetworkOfferingId()).thenReturn(8L);
        lenient().when(networkNat.getState()).thenReturn(Network.State.Implemented);
        lenient().when(networkNat.getGuestType()).thenReturn(Network.GuestType.Isolated);
        when(networkNat.getId()).thenReturn(3L);
        lenient().when(networkNat.getVpcId()).thenReturn(null);
        when(networkNat.getId()).thenReturn(3L);

        IPAddressVO sourceNat = new IPAddressVO(new Ip("192.0.0.2"), 1L, 1L, 1L,true);

        Mockito.lenient().when(networkDao.findById(3L)).thenReturn(networkNat);
        doReturn(sourceNat).when(ipAddressManager).getExistingSourceNatInNetwork(1L, 3L);

        boolean isSourceNat = ipAddressManager.isSourceNatAvailableForNetwork(account, ipAddressVO, networkNat);

        assertFalse("Source NAT should be false", isSourceNat);
    }

    @Test
    public void isIpEqualsGatewayOrNetworkOfferingsEmptyTestNetworkOfferingsEmptyAndCidrNull() {
        Network network = setTestIsIpEqualsGatewayOrNetworkOfferingsEmpty(0l, "gateway", "ip6Gateway", null, new ArrayList<Service>());

        boolean result = ipAddressManager.isIpEqualsGatewayOrNetworkOfferingsEmpty(network, "requestedIp");

        Mockito.verify(networkModel).listNetworkOfferingServices(Mockito.anyLong());
        Assert.assertTrue(result);
    }

    private Network setTestIsIpEqualsGatewayOrNetworkOfferingsEmpty(long networkOfferingId, String gateway, String ip6Gateway, String cidr, List<Service> services) {
        Network network = mock(Network.class);
        Mockito.when(network.getNetworkOfferingId()).thenReturn(networkOfferingId);
        Mockito.when(network.getGateway()).thenReturn(gateway);
        Mockito.when(network.getIp6Gateway()).thenReturn(ip6Gateway);
        Mockito.when(network.getCidr()).thenReturn(cidr);
        Mockito.doReturn(services).when(networkModel).listNetworkOfferingServices(Mockito.anyLong());
        return network;
    }

    @Test
    public void isPublicIpAddressStillInQuarantineTestRemovedDateIsNullAndCurrentDateIsEqualToEndDateShouldReturnFalse() {
        Date endDate = currentDate;

        Mockito.when(publicIpQuarantineVOMock.getRemoved()).thenReturn(null);
        Mockito.when(publicIpQuarantineVOMock.getEndDate()).thenReturn(endDate);

        boolean result = ipAddressManager.isPublicIpAddressStillInQuarantine(publicIpQuarantineVOMock, currentDate);

        Assert.assertFalse(result);
    }

    @Test
    public void isPublicIpAddressStillInQuarantineTestRemovedDateIsNullAndEndDateIsBeforeCurrentDateShouldReturnFalse() {
        Date endDate = beforeCurrentDate;

        Mockito.when(publicIpQuarantineVOMock.getRemoved()).thenReturn(null);
        Mockito.when(publicIpQuarantineVOMock.getEndDate()).thenReturn(endDate);

        boolean result = ipAddressManager.isPublicIpAddressStillInQuarantine(publicIpQuarantineVOMock, currentDate);

        Assert.assertFalse(result);
    }

    @Test
    public void isPublicIpAddressStillInQuarantineTestRemovedDateIsNullAndEndDateIsAfterCurrentDateShouldReturnTrue() {
        Date endDate = afterCurrentDate;

        Mockito.when(publicIpQuarantineVOMock.getRemoved()).thenReturn(null);
        Mockito.when(publicIpQuarantineVOMock.getEndDate()).thenReturn(endDate);

        boolean result = ipAddressManager.isPublicIpAddressStillInQuarantine(publicIpQuarantineVOMock, currentDate);

        Assert.assertTrue(result);
    }

    @Test
    public void isPublicIpAddressStillInQuarantineTestRemovedDateIsEqualCurrentDateShouldReturnFalse() {
        Date removedDate = currentDate;

        Mockito.when(publicIpQuarantineVOMock.getEndDate()).thenReturn(currentDate);
        Mockito.when(publicIpQuarantineVOMock.getRemoved()).thenReturn(removedDate);

        boolean result = ipAddressManager.isPublicIpAddressStillInQuarantine(publicIpQuarantineVOMock, currentDate);

        Assert.assertFalse(result);
    }

    @Test
    public void isPublicIpAddressStillInQuarantineTestRemovedDateIsBeforeCurrentDateShouldReturnFalse() {
        Date removedDate = beforeCurrentDate;

        Mockito.when(publicIpQuarantineVOMock.getRemoved()).thenReturn(removedDate);
        Mockito.when(publicIpQuarantineVOMock.getEndDate()).thenReturn(null);

        boolean result = ipAddressManager.isPublicIpAddressStillInQuarantine(publicIpQuarantineVOMock, currentDate);

        Assert.assertFalse(result);
    }

    @Test
    public void isPublicIpAddressStillInQuarantineTestRemovedDateIsAfterCurrentDateShouldReturnTrue() {
        Date removedDate = afterCurrentDate;

        Mockito.when(publicIpQuarantineVOMock.getRemoved()).thenReturn(removedDate);
        Mockito.when(publicIpQuarantineVOMock.getEndDate()).thenReturn(null);

        boolean result = ipAddressManager.isPublicIpAddressStillInQuarantine(publicIpQuarantineVOMock, currentDate);

        Assert.assertTrue(result);
    }

    @Test
    public void checkIfPublicIpAddressIsNotInQuarantineAndCanBeAllocatedTestIpIsNotInQuarantineShouldReturnTrue() {
        Mockito.when(ipAddressMock.getId()).thenReturn(dummyID);
        Mockito.when(publicIpQuarantineDaoMock.findByPublicIpAddressId(Mockito.anyLong())).thenReturn(null);

        boolean result = ipAddressManager.canPublicIpAddressBeAllocated(ipAddressMock, account);

        Assert.assertTrue(result);
    }

    @Test
    public void checkIfPublicIpAddressIsNotInQuarantineAndCanBeAllocatedTestIpIsNoLongerInQuarantineShouldReturnTrue() {
        Mockito.when(ipAddressMock.getId()).thenReturn(dummyID);
        Mockito.when(publicIpQuarantineDaoMock.findByPublicIpAddressId(Mockito.anyLong())).thenReturn(publicIpQuarantineVOMock);
        Mockito.doReturn(false).when(ipAddressManager).isPublicIpAddressStillInQuarantine(Mockito.any(PublicIpQuarantineVO.class), Mockito.any(Date.class));

        boolean result = ipAddressManager.canPublicIpAddressBeAllocated(ipAddressMock, newOwnerMock);

        Assert.assertTrue(result);
    }

    @Test
    public void checkIfPublicIpAddressIsNotInQuarantineAndCanBeAllocatedTestIpIsInQuarantineAndThePreviousOwnerIsTheSameAsTheNewOwnerShouldReturnTrue() {
        Mockito.when(ipAddressMock.getId()).thenReturn(dummyID);
        Mockito.when(publicIpQuarantineDaoMock.findByPublicIpAddressId(Mockito.anyLong())).thenReturn(publicIpQuarantineVOMock);

        Mockito.doReturn(true).when(ipAddressManager).isPublicIpAddressStillInQuarantine(Mockito.any(PublicIpQuarantineVO.class), Mockito.any(Date.class));
        Mockito.doNothing().when(ipAddressManager).removePublicIpAddressFromQuarantine(Mockito.anyLong(), Mockito.anyString());

        Mockito.when(publicIpQuarantineVOMock.getPreviousOwnerId()).thenReturn(dummyID);
        Mockito.when(accountManagerMock.getAccount(Mockito.anyLong())).thenReturn(previousOwnerMock);
        Mockito.when(previousOwnerMock.getUuid()).thenReturn(UUID);
        Mockito.when(newOwnerMock.getUuid()).thenReturn(UUID);

        boolean result = ipAddressManager.canPublicIpAddressBeAllocated(ipAddressMock, newOwnerMock);

        Assert.assertTrue(result);
    }

    @Test
    public void checkIfPublicIpAddressIsNotInQuarantineAndCanBeAllocatedTestIpIsInQuarantineAndThePreviousOwnerIsDifferentFromTheNewOwnerShouldReturnFalse() {
        final String UUID_2 = "uuid_2";

        Mockito.when(ipAddressMock.getId()).thenReturn(dummyID);
        Mockito.when(publicIpQuarantineDaoMock.findByPublicIpAddressId(Mockito.anyLong())).thenReturn(publicIpQuarantineVOMock);

        Mockito.doReturn(true).when(ipAddressManager).isPublicIpAddressStillInQuarantine(Mockito.any(PublicIpQuarantineVO.class), Mockito.any(Date.class));

        Mockito.when(publicIpQuarantineVOMock.getPreviousOwnerId()).thenReturn(dummyID);
        Mockito.when(accountManagerMock.getAccount(Mockito.anyLong())).thenReturn(previousOwnerMock);
        Mockito.when(previousOwnerMock.getUuid()).thenReturn(UUID);
        Mockito.when(newOwnerMock.getUuid()).thenReturn(UUID_2);

        boolean result = ipAddressManager.canPublicIpAddressBeAllocated(ipAddressMock, newOwnerMock);

        Assert.assertFalse(result);
    }

    @Test
    public void removePublicIpAddressFromQuarantineTestPersistsObject() {
        Long quarantineProcessId = 100L;
        Long publicAddressId = 200L;
        Long callingAccountId = 300L;
        String removalReason = "removalReason";

        try (MockedStatic<CallContext> callContextMockedStatic = Mockito.mockStatic(CallContext.class)) {
            Mockito.doReturn(publicIpQuarantineVOMock).when(publicIpQuarantineDaoMock).findById(quarantineProcessId);
            Mockito.when(publicIpQuarantineVOMock.getPublicIpAddressId()).thenReturn(publicAddressId);
            Mockito.doReturn(ipAddressVO).when(ipAddressDao).findById(publicAddressId);

            CallContext callContextMock = Mockito.mock(CallContext.class);
            Mockito.when(callContextMock.getCallingAccountId()).thenReturn(callingAccountId);
            callContextMockedStatic.when(CallContext::current).thenReturn(callContextMock);

            ipAddressManager.removePublicIpAddressFromQuarantine(quarantineProcessId, removalReason);

            Mockito.verify(publicIpQuarantineVOMock).setRemoved(Mockito.any());
            Mockito.verify(publicIpQuarantineVOMock).setRemovalReason(removalReason);
            Mockito.verify(publicIpQuarantineVOMock).setRemoverAccountId(callingAccountId);
            Mockito.verify(publicIpQuarantineDaoMock).persist(publicIpQuarantineVOMock);
        }
    }

    @Test
    public void updateSourceNatIpAddress() throws Exception {
        IPAddressVO requestedIp = Mockito.mock(IPAddressVO.class);
        IPAddressVO oldIp = Mockito.mock(IPAddressVO.class);
        List<IPAddressVO> userIps = new Vector<>();
        userIps.add(oldIp);
        ipAddressManager.updateSourceNatIpAddress(requestedIp, userIps);
        verify(requestedIp).setSourceNat(true);
        verify(oldIp).setSourceNat(false);
    }

    private void prepareForCheckIfIpResourceCountShouldBeUpdatedTests() {
        Mockito.when(ipAddressVoMock.getAssociatedWithNetworkId()).thenReturn(1L);
        Mockito.when(ipAddressVoMock.getVpcId()).thenReturn(1L);
        doReturn(false).when(ipAddressManager).isIpDedicated(Mockito.any());
        Mockito.when(ipAddressVoMock.getState()).thenReturn(IpAddress.State.Allocating);
    }

    @Test
    public void checkIfIpResourceCountShouldBeUpdatedTestIpIsDirectReturnFalse() {
        prepareForCheckIfIpResourceCountShouldBeUpdatedTests();
        Mockito.when(ipAddressVoMock.getAssociatedWithNetworkId()).thenReturn(null);
        Mockito.when(ipAddressVoMock.getVpcId()).thenReturn(null);

        boolean result = ipAddressManager.checkIfIpResourceCountShouldBeUpdated(ipAddressVoMock);

        Assert.assertFalse(result);
    }

    @Test
    public void checkIfIpResourceCountShouldBeUpdatedTestIpIsDedicatedReturnFalse() {
        prepareForCheckIfIpResourceCountShouldBeUpdatedTests();
        doReturn(true).when(ipAddressManager).isIpDedicated(Mockito.any());

        boolean result = ipAddressManager.checkIfIpResourceCountShouldBeUpdated(ipAddressVoMock);

        Assert.assertFalse(result);
    }

    @Test
    public void checkIfIpResourceCountShouldBeUpdatedTestIpIsReservedReturnFalse() {
        prepareForCheckIfIpResourceCountShouldBeUpdatedTests();
        Mockito.when(ipAddressVoMock.getState()).thenReturn(IpAddress.State.Reserved);

        boolean result = ipAddressManager.checkIfIpResourceCountShouldBeUpdated(ipAddressVoMock);

        Assert.assertFalse(result);
    }

    @Test
    public void checkIfIpResourceCountShouldBeUpdatedTestIpIsAssociatedToNetworkAndNotDedicatedAndNotReservedReturnTrue() {
        prepareForCheckIfIpResourceCountShouldBeUpdatedTests();

        boolean result = ipAddressManager.checkIfIpResourceCountShouldBeUpdated(ipAddressVoMock);

        Assert.assertTrue(result);
    }

    @Test
    public void checkIfIpResourceCountShouldBeUpdatedTestIpIsAssociatedToVpcAndNotDedicatedAndNotReservedReturnTrue() {
        prepareForCheckIfIpResourceCountShouldBeUpdatedTests();
        Mockito.when(ipAddressVoMock.getAssociatedWithNetworkId()).thenReturn(null);

        boolean result = ipAddressManager.checkIfIpResourceCountShouldBeUpdated(ipAddressVoMock);

        Assert.assertTrue(result);
    }


    private FirewallRule makeRule(Long networkId, Long vpcId, FirewallRule.Purpose purpose,
            FirewallRule.TrafficType trafficType) {
        FirewallRule rule = mock(FirewallRule.class);
        lenient().when(rule.getNetworkId()).thenReturn(networkId);
        lenient().when(rule.getVpcId()).thenReturn(vpcId);
        lenient().when(rule.getPurpose()).thenReturn(purpose);
        lenient().when(rule.getTrafficType()).thenReturn(trafficType);
        return rule;
    }

    /** Stub the two IP-association helper methods so they are no-ops. */
    private void stubIpAssocHelpers() throws ResourceUnavailableException {
        doReturn(false).when(ipAddressManager).checkIfIpAssocRequired(any(Network.class), anyBoolean(), any());
    }

    /**
     * Test: Non-VPC rules still resolve via networkId (backward compatibility).
     */
    @Test
    public void applyRulesNonVpcRuleStillWorksViaNetworkId() throws ResourceUnavailableException {
        long networkId = 10L;
        NetworkVO network = mock(NetworkVO.class);
        when(network.getId()).thenReturn(networkId);
        when(networkDao.findById(networkId)).thenReturn(network);

        FirewallRule rule = makeRule(networkId, null, FirewallRule.Purpose.Firewall, FirewallRule.TrafficType.Ingress);
        NetworkRuleApplier applier = mock(NetworkRuleApplier.class);

        when(ipAddressDao.listByAssociatedNetwork(networkId, null)).thenReturn(new ArrayList<>());
        stubIpAssocHelpers();

        boolean result = ipAddressManager.applyRules(
                Collections.singletonList(rule), FirewallRule.Purpose.Firewall, applier, false);

        assertTrue(result);
        verify(networkDao).findById(networkId);
        verify(applier).applyRules(network, FirewallRule.Purpose.Firewall, Collections.singletonList(rule));
    }

    /**
     * Test: VPC rule resolves network via VpcManager.getVpcNetworks()
     * when networkId is null but vpcId is set.
     */
    @Test
    public void applyRulesVpcRuleResolvesNetworkViaVpcManager() throws ResourceUnavailableException {
        long vpcId = 20L;
        long resolvedNetworkId = 30L;
        NetworkVO resolvedNetwork = mock(NetworkVO.class);
        when(resolvedNetwork.getId()).thenReturn(resolvedNetworkId);
        when(networkDao.findById(resolvedNetworkId)).thenReturn(resolvedNetwork);

        doReturn(Collections.singletonList(resolvedNetwork)).when(vpcMgr).getVpcNetworks(vpcId);

        FirewallRule rule = makeRule(null, vpcId, FirewallRule.Purpose.Firewall, FirewallRule.TrafficType.Ingress);
        NetworkRuleApplier applier = mock(NetworkRuleApplier.class);

        IPAddressVO vpcIp = mock(IPAddressVO.class);
        when(vpcIp.getVlanId()).thenReturn(1L);
        VlanVO vlan = mock(VlanVO.class);
        when(ipAddressDao.listByAssociatedVpc(vpcId, null)).thenReturn(Collections.singletonList(vpcIp));
        when(vlanDao.findById(1L)).thenReturn(vlan);

        stubIpAssocHelpers();

        boolean result = ipAddressManager.applyRules(
                Collections.singletonList(rule), FirewallRule.Purpose.Firewall, applier, false);

        assertTrue(result);
        verify(vpcMgr).getVpcNetworks(vpcId);
        verify(ipAddressDao).listByAssociatedVpc(vpcId, null);
        verify(applier).applyRules(resolvedNetwork, FirewallRule.Purpose.Firewall, Collections.singletonList(rule));
    }


    /**
     * Test: For VPC egress firewall rules, IP collection should be skipped.
     */
    @Test
    public void applyRulesVpcEgressFirewallRuleSkipsIpCollection() throws ResourceUnavailableException {
        long vpcId = 20L;
        long resolvedNetworkId = 30L;
        NetworkVO resolvedNetwork = mock(NetworkVO.class);
        when(resolvedNetwork.getId()).thenReturn(resolvedNetworkId);
        when(networkDao.findById(resolvedNetworkId)).thenReturn(resolvedNetwork);
        doReturn(Collections.singletonList(resolvedNetwork)).when(vpcMgr).getVpcNetworks(vpcId);

        FirewallRule rule = makeRule(null, vpcId, FirewallRule.Purpose.Firewall, FirewallRule.TrafficType.Egress);
        NetworkRuleApplier applier = mock(NetworkRuleApplier.class);

        stubIpAssocHelpers();

        boolean result = ipAddressManager.applyRules(
                Collections.singletonList(rule), FirewallRule.Purpose.Firewall, applier, false);

        assertTrue(result);
        verify(ipAddressDao, never()).listByAssociatedVpc(anyLong(), any());
        verify(applier).applyRules(resolvedNetwork, FirewallRule.Purpose.Firewall, Collections.singletonList(rule));
    }

    /**
     * Test: VPC ingress firewall rules collect public IPs from VPC (listByAssociatedVpc),
     * NOT from network (listByAssociatedNetwork).
     */
    @Test
    public void applyRulesVpcIngressRuleCollectsIpsFromVpcNotNetwork() throws ResourceUnavailableException {
        long vpcId = 20L;
        long resolvedNetworkId = 30L;
        NetworkVO resolvedNetwork = mock(NetworkVO.class);
        when(resolvedNetwork.getId()).thenReturn(resolvedNetworkId);
        when(networkDao.findById(resolvedNetworkId)).thenReturn(resolvedNetwork);
        doReturn(Collections.singletonList(resolvedNetwork)).when(vpcMgr).getVpcNetworks(vpcId);

        IPAddressVO vpcIp = mock(IPAddressVO.class);
        when(vpcIp.getVlanId()).thenReturn(1L);
        VlanVO vlan = mock(VlanVO.class);
        when(ipAddressDao.listByAssociatedVpc(vpcId, null)).thenReturn(Collections.singletonList(vpcIp));
        when(vlanDao.findById(1L)).thenReturn(vlan);

        stubIpAssocHelpers();

        NetworkRuleApplier applier = mock(NetworkRuleApplier.class);
        FirewallRule rule = makeRule(null, vpcId, FirewallRule.Purpose.Firewall, FirewallRule.TrafficType.Ingress);

        ipAddressManager.applyRules(Collections.singletonList(rule), FirewallRule.Purpose.Firewall, applier, false);

        verify(ipAddressDao).listByAssociatedVpc(vpcId, null);
        verify(ipAddressDao, never()).listByAssociatedNetwork(anyLong(), any());
    }

    /**
     * Test: Error handling respects continueOnError flag.
     * When continueOnError=true, exceptions are caught and false is returned.
     */
    @Test
    public void applyRulesVpcRuleErrorHandlingWithContinueOnErrorTrue() throws ResourceUnavailableException {
        long vpcId = 20L;
        long resolvedNetworkId = 30L;
        NetworkVO resolvedNetwork = mock(NetworkVO.class);
        when(resolvedNetwork.getId()).thenReturn(resolvedNetworkId);
        when(networkDao.findById(resolvedNetworkId)).thenReturn(resolvedNetwork);
        doReturn(Collections.singletonList(resolvedNetwork)).when(vpcMgr).getVpcNetworks(vpcId);

        IPAddressVO vpcIp = mock(IPAddressVO.class);
        when(vpcIp.getVlanId()).thenReturn(1L);
        VlanVO vlan = mock(VlanVO.class);
        when(ipAddressDao.listByAssociatedVpc(vpcId, null)).thenReturn(Collections.singletonList(vpcIp));
        when(vlanDao.findById(1L)).thenReturn(vlan);

        stubIpAssocHelpers();

        NetworkRuleApplier applier = mock(NetworkRuleApplier.class);
        when(applier.applyRules(any(), any(), any())).thenThrow(new ResourceUnavailableException("test", Network.class, 0L));

        FirewallRule rule = makeRule(null, vpcId, FirewallRule.Purpose.Firewall, FirewallRule.TrafficType.Ingress);

        boolean result = ipAddressManager.applyRules(
                Collections.singletonList(rule), FirewallRule.Purpose.Firewall, applier, true);

        assertFalse(result);
    }
}