Merge pull request #1360 from greenqloud/pr-refactor-system-vm-network-creation

Refactor system VM default network creationTwo small commits which moves the retrieval of the default network for the console proxy and the SSVM into a separate protected method. It's a small change that makes the code more readable/maintainable and also makes the class more suitable for overriding should one want to do this. It's forward-ported from our 4.2 branch.

No new tests since this should not change any functionality, and thus should be covered by the existing unit tests.

Now on the master branch (#1359 was on the wrong branch).

* pr/1360:
  Refactor ssvm default network retrieval.
  Refactor console proxy default network retrieval.

Signed-off-by: Will Stevens <williamstevens@gmail.com>

server/src/com/cloud/consoleproxy/ConsoleProxyManagerImpl.java

@@ -41,6 +41,7 @@ import org.apache.cloudstack.storage.datastore.db.PrimaryDataStoreDao;
 import org.apache.cloudstack.storage.datastore.db.StoragePoolVO;
 import org.apache.cloudstack.storage.datastore.db.TemplateDataStoreDao;
 import org.apache.cloudstack.storage.datastore.db.TemplateDataStoreVO;
+import org.apache.commons.collections.CollectionUtils;
 import org.apache.log4j.Logger;
 
 import com.cloud.agent.AgentManager;
@@ -666,6 +667,78 @@ public class ConsoleProxyManagerImpl extends ManagerBase implements ConsoleProxy
         return null;
     }
 
+    /**
+     * Get the default network for the console proxy VM, based on the zone it is in. Delegates to
+     * either {@link #getDefaultNetworkForZone(DataCenter)} or {@link #getDefaultNetworkForAdvancedSGZone(DataCenter)},
+     * depending on the zone network type and whether or not security groups are enabled in the zone.
+     * @param dc - The zone (DataCenter) of the console proxy VM.
+     * @return The default network for use with the console proxy VM.
+     */
+    protected NetworkVO getDefaultNetworkForCreation(DataCenter dc) {
+        if (dc.getNetworkType() == NetworkType.Advanced) {
+            return getDefaultNetworkForAdvancedZone(dc);
+        } else {
+            return getDefaultNetworkForBasicZone(dc);
+        }
+    }
+
+    /**
+     * Get default network for a console proxy VM starting up in an advanced zone. If the zone
+     * is security group-enabled, the first network found that supports SG services is returned.
+     * If the zone is not SG-enabled, the Public network is returned.
+     * @param dc - The zone.
+     * @return The selected default network.
+     * @throws CloudRuntimeException - If the zone is not a valid choice or a network couldn't be found.
+     */
+    protected NetworkVO getDefaultNetworkForAdvancedZone(DataCenter dc) {
+        if (dc.getNetworkType() != NetworkType.Advanced) {
+            throw new CloudRuntimeException("Zone " + dc + " is not advanced.");
+        }
+
+        if (dc.isSecurityGroupEnabled()) {
+            List<NetworkVO> networks = _networkDao.listByZoneSecurityGroup(dc.getId());
+            if (CollectionUtils.isEmpty(networks)) {
+                throw new CloudRuntimeException("Can not found security enabled network in SG Zone " + dc);
+            }
+
+            return networks.get(0);
+        }
+        else {
+            TrafficType defaultTrafficType = TrafficType.Public;
+            List<NetworkVO> defaultNetworks = _networkDao.listByZoneAndTrafficType(dc.getId(), defaultTrafficType);
+
+            // api should never allow this situation to happen
+            if (defaultNetworks.size() != 1) {
+                throw new CloudRuntimeException("Found " + defaultNetworks.size() + " networks of type " + defaultTrafficType + " when expect to find 1");
+            }
+
+            return defaultNetworks.get(0);
+        }
+    }
+
+    /**
+     * Get default network for console proxy VM for starting up in a basic zone. Basic zones select
+     * the Guest network whether or not the zone is SG-enabled.
+     * @param dc - The zone.
+     * @return The default network according to the zone's network selection rules.
+     * @throws CloudRuntimeException - If the zone is not a valid choice or a network couldn't be found.
+     */
+    protected NetworkVO getDefaultNetworkForBasicZone(DataCenter dc) {
+        if (dc.getNetworkType() != NetworkType.Basic) {
+            throw new CloudRuntimeException("Zone " + dc + "is not basic.");
+        }
+
+        TrafficType defaultTrafficType = TrafficType.Guest;
+        List<NetworkVO> defaultNetworks = _networkDao.listByZoneAndTrafficType(dc.getId(), defaultTrafficType);
+
+        // api should never allow this situation to happen
+        if (defaultNetworks.size() != 1) {
+            throw new CloudRuntimeException("Found " + defaultNetworks.size() + " networks of type " + defaultTrafficType + " when expect to find 1");
+        }
+
+        return defaultNetworks.get(0);
+    }
+
     protected Map<String, Object> createProxyInstance(long dataCenterId, VMTemplateVO template) throws ConcurrentOperationException {
 
         long id = _consoleProxyDao.getNextInSequence(Long.class, "id");
@@ -675,26 +748,7 @@ public class ConsoleProxyManagerImpl extends ManagerBase implements ConsoleProxy
 
         DataCenterDeployment plan = new DataCenterDeployment(dataCenterId);
 
-        NetworkVO defaultNetwork = null;
-        if (dc.getNetworkType() == NetworkType.Advanced && dc.isSecurityGroupEnabled()) {
-            List<NetworkVO> networks = _networkDao.listByZoneSecurityGroup(dataCenterId);
-            if (networks == null || networks.size() == 0) {
-                throw new CloudRuntimeException("Can not found security enabled network in SG Zone " + dc);
-            }
-            defaultNetwork = networks.get(0);
-        } else {
-            TrafficType defaultTrafficType = TrafficType.Public;
-            if (dc.getNetworkType() == NetworkType.Basic || dc.isSecurityGroupEnabled()) {
-                defaultTrafficType = TrafficType.Guest;
-            }
-            List<NetworkVO> defaultNetworks = _networkDao.listByZoneAndTrafficType(dataCenterId, defaultTrafficType);
-
-            // api should never allow this situation to happen
-            if (defaultNetworks.size() != 1) {
-                throw new CloudRuntimeException("Found " + defaultNetworks.size() + " networks of type " + defaultTrafficType + " when expect to find 1");
-            }
-            defaultNetwork = defaultNetworks.get(0);
-        }
+        NetworkVO defaultNetwork = getDefaultNetworkForCreation(dc);
 
         List<? extends NetworkOffering> offerings =
             _networkModel.getSystemAccountNetworkOfferings(NetworkOffering.SystemControlNetwork, NetworkOffering.SystemManagementNetwork);

server/test/com/cloud/consoleproxy/ConsoleProxyManagerTest.java

@@ -17,7 +17,17 @@
 
 package com.cloud.consoleproxy;
 
+import static org.mockito.AdditionalMatchers.not;
+import static org.mockito.Matchers.any;
+import static org.mockito.Matchers.anyLong;
+import static org.mockito.Matchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import java.util.Collections;
+
 import org.apache.log4j.Logger;
+import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 import org.mockito.Mock;
@@ -25,7 +35,15 @@ import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
 import org.springframework.test.util.ReflectionTestUtils;
 
+import com.cloud.dc.DataCenterVO;
+import com.cloud.dc.DataCenter;
+import com.cloud.dc.DataCenter.NetworkType;
+import com.cloud.dc.dao.DataCenterDao;
+import com.cloud.network.Networks.TrafficType;
+import com.cloud.network.dao.NetworkDao;
+import com.cloud.network.dao.NetworkVO;
 import com.cloud.utils.db.GlobalLock;
+import com.cloud.utils.exception.CloudRuntimeException;
 import com.cloud.vm.ConsoleProxyVO;
 
 public class ConsoleProxyManagerTest {
@@ -37,13 +55,22 @@ public class ConsoleProxyManagerTest {
     @Mock
     ConsoleProxyVO proxyVO;
     @Mock
+    DataCenterDao _dcDao;
+    @Mock
+    NetworkDao _networkDao;
+    @Mock
     ConsoleProxyManagerImpl cpvmManager;
 
     @Before
     public void setup() throws Exception {
         MockitoAnnotations.initMocks(this);
         ReflectionTestUtils.setField(cpvmManager, "_allocProxyLock", globalLock);
+        ReflectionTestUtils.setField(cpvmManager, "_dcDao", _dcDao);
+        ReflectionTestUtils.setField(cpvmManager, "_networkDao", _networkDao);
         Mockito.doCallRealMethod().when(cpvmManager).expandPool(Mockito.anyLong(), Mockito.anyObject());
+        Mockito.doCallRealMethod().when(cpvmManager).getDefaultNetworkForCreation(Mockito.any(DataCenter.class));
+        Mockito.doCallRealMethod().when(cpvmManager).getDefaultNetworkForAdvancedZone(Mockito.any(DataCenter.class));
+        Mockito.doCallRealMethod().when(cpvmManager).getDefaultNetworkForBasicZone(Mockito.any(DataCenter.class));
     }
 
     @Test
@@ -87,4 +114,136 @@ public class ConsoleProxyManagerTest {
 
         cpvmManager.expandPool(new Long(1), new Object());
     }
+
+    @Test
+    public void getDefaultNetworkForAdvancedNonSG() {
+        DataCenterVO dc = mock(DataCenterVO.class);
+        when(dc.getNetworkType()).thenReturn(NetworkType.Advanced);
+        when(dc.isSecurityGroupEnabled()).thenReturn(false);
+
+        when(_dcDao.findById(Mockito.anyLong())).thenReturn(dc);
+
+        NetworkVO network = Mockito.mock(NetworkVO.class);
+        NetworkVO badNetwork = Mockito.mock(NetworkVO.class);
+        when(_networkDao.listByZoneAndTrafficType(anyLong(), eq(TrafficType.Public)))
+                    .thenReturn(Collections.singletonList(network));
+
+        when(_networkDao.listByZoneAndTrafficType(anyLong(), not(eq(TrafficType.Public))))
+        .thenReturn(Collections.singletonList(badNetwork));
+
+        when(_networkDao.listByZoneSecurityGroup(anyLong()))
+                    .thenReturn(Collections.singletonList(badNetwork));
+
+        NetworkVO returnedNetwork = cpvmManager.getDefaultNetworkForAdvancedZone(dc);
+
+        Assert.assertNotNull(returnedNetwork);
+        Assert.assertEquals(network, returnedNetwork);
+        Assert.assertNotEquals(badNetwork, returnedNetwork);
+    }
+
+    @Test
+    public void getDefaultNetworkForAdvancedSG() {
+        DataCenterVO dc = Mockito.mock(DataCenterVO.class);
+        when(dc.getNetworkType()).thenReturn(NetworkType.Advanced);
+        when(dc.isSecurityGroupEnabled()).thenReturn(true);
+
+        when(_dcDao.findById(Mockito.anyLong())).thenReturn(dc);
+
+        NetworkVO network = Mockito.mock(NetworkVO.class);
+        NetworkVO badNetwork = Mockito.mock(NetworkVO.class);
+        when(_networkDao.listByZoneAndTrafficType(anyLong(), any(TrafficType.class)))
+                    .thenReturn(Collections.singletonList(badNetwork));
+
+        when(_networkDao.listByZoneSecurityGroup(anyLong()))
+                    .thenReturn(Collections.singletonList(network));
+
+        NetworkVO returnedNetwork = cpvmManager.getDefaultNetworkForAdvancedZone(dc);
+
+        Assert.assertEquals(network, returnedNetwork);
+        Assert.assertNotEquals(badNetwork, returnedNetwork);
+    }
+
+    @Test
+    public void getDefaultNetworkForBasicNonSG() {
+        DataCenterVO dc = Mockito.mock(DataCenterVO.class);
+        when(dc.getNetworkType()).thenReturn(NetworkType.Basic);
+        when(dc.isSecurityGroupEnabled()).thenReturn(false);
+
+        when(_dcDao.findById(Mockito.anyLong())).thenReturn(dc);
+
+        NetworkVO network = Mockito.mock(NetworkVO.class);
+        NetworkVO badNetwork = Mockito.mock(NetworkVO.class);
+        when(_networkDao.listByZoneAndTrafficType(anyLong(), eq(TrafficType.Guest)))
+                    .thenReturn(Collections.singletonList(network));
+
+        when(_networkDao.listByZoneAndTrafficType(anyLong(), not(eq(TrafficType.Guest))))
+                    .thenReturn(Collections.singletonList(badNetwork));
+
+        NetworkVO returnedNetwork = cpvmManager.getDefaultNetworkForBasicZone(dc);
+        Assert.assertNotNull(returnedNetwork);
+        Assert.assertEquals(network, returnedNetwork);
+        Assert.assertNotEquals(badNetwork, returnedNetwork);
+    }
+
+    @Test
+    public void getDefaultNetworkForBasicSG() {
+        DataCenterVO dc = Mockito.mock(DataCenterVO.class);
+        when(dc.getNetworkType()).thenReturn(NetworkType.Basic);
+        when(dc.isSecurityGroupEnabled()).thenReturn(true);
+
+        when(_dcDao.findById(Mockito.anyLong())).thenReturn(dc);
+
+        NetworkVO network = Mockito.mock(NetworkVO.class);
+        NetworkVO badNetwork = Mockito.mock(NetworkVO.class);
+        when(_networkDao.listByZoneAndTrafficType(anyLong(), eq(TrafficType.Guest)))
+                    .thenReturn(Collections.singletonList(network));
+
+        when(_networkDao.listByZoneAndTrafficType(anyLong(), not(eq(TrafficType.Guest))))
+                    .thenReturn(Collections.singletonList(badNetwork));
+
+        NetworkVO returnedNetwork = cpvmManager.getDefaultNetworkForBasicZone(dc);
+
+        Assert.assertNotNull(returnedNetwork);
+        Assert.assertEquals(network, returnedNetwork);
+        Assert.assertNotEquals(badNetwork, returnedNetwork);
+    }
+
+    //also test invalid input
+    @Test(expected=CloudRuntimeException.class)
+    public void getDefaultNetworkForBasicSGWrongZoneType() {
+        DataCenterVO dc = Mockito.mock(DataCenterVO.class);
+        when(dc.getNetworkType()).thenReturn(NetworkType.Advanced);
+        when(dc.isSecurityGroupEnabled()).thenReturn(true);
+
+        when(_dcDao.findById(Mockito.anyLong())).thenReturn(dc);
+
+        NetworkVO network = Mockito.mock(NetworkVO.class);
+        NetworkVO badNetwork = Mockito.mock(NetworkVO.class);
+        when(_networkDao.listByZoneAndTrafficType(anyLong(), eq(TrafficType.Guest)))
+                    .thenReturn(Collections.singletonList(network));
+
+        when(_networkDao.listByZoneAndTrafficType(anyLong(), not(eq(TrafficType.Guest))))
+                    .thenReturn(Collections.singletonList(badNetwork));
+
+        cpvmManager.getDefaultNetworkForBasicZone(dc);
+    }
+
+    @Test(expected=CloudRuntimeException.class)
+    public void getDefaultNetworkForAdvancedWrongZoneType() {
+        DataCenterVO dc = Mockito.mock(DataCenterVO.class);
+        when(dc.getNetworkType()).thenReturn(NetworkType.Basic);
+        when(dc.isSecurityGroupEnabled()).thenReturn(true);
+
+        when(_dcDao.findById(Mockito.anyLong())).thenReturn(dc);
+
+        NetworkVO network = Mockito.mock(NetworkVO.class);
+        NetworkVO badNetwork = Mockito.mock(NetworkVO.class);
+        when(_networkDao.listByZoneAndTrafficType(anyLong(), any(TrafficType.class)))
+                    .thenReturn(Collections.singletonList(badNetwork));
+
+        when(_networkDao.listByZoneSecurityGroup(anyLong()))
+                    .thenReturn(Collections.singletonList(network));
+
+        cpvmManager.getDefaultNetworkForAdvancedZone(dc);
+    }
 }

services/secondary-storage/controller/src/org/apache/cloudstack/secondarystorage/SecondaryStorageManagerImpl.java

@@ -43,6 +43,7 @@ import org.apache.cloudstack.storage.datastore.db.ImageStoreVO;
 import org.apache.cloudstack.storage.datastore.db.TemplateDataStoreDao;
 import org.apache.cloudstack.storage.datastore.db.VolumeDataStoreDao;
 import org.apache.cloudstack.utils.identity.ManagementServerNode;
+import org.apache.commons.collections.CollectionUtils;
 import org.apache.log4j.Logger;
 
 import com.cloud.agent.AgentManager;
@@ -521,6 +522,76 @@ public class SecondaryStorageManagerImpl extends ManagerBase implements Secondar
         return null;
     }
 
+    /**
+     * Get the default network for the secondary storage VM, based on the zone it is in. Delegates to
+     * either {@link #getDefaultNetworkForZone(DataCenter)} or {@link #getDefaultNetworkForAdvancedSGZone(DataCenter)},
+     * depending on the zone network type and whether or not security groups are enabled in the zone.
+     * @param dc - The zone (DataCenter) of the secondary storage VM.
+     * @return The default network for use with the secondary storage VM.
+     */
+    protected NetworkVO getDefaultNetworkForCreation(DataCenter dc) {
+        if (dc.getNetworkType() == NetworkType.Advanced) {
+            return getDefaultNetworkForAdvancedZone(dc);
+        } else {
+            return getDefaultNetworkForBasicZone(dc);
+        }
+    }
+
+    /**
+     * Get default network for a secondary storage VM starting up in an advanced zone. If the zone
+     * is security group-enabled, the first network found that supports SG services is returned.
+     * If the zone is not SG-enabled, the Public network is returned.
+     * @param dc - The zone.
+     * @return The selected default network.
+     * @throws CloudRuntimeException - If the zone is not a valid choice or a network couldn't be found.
+     */
+    protected NetworkVO getDefaultNetworkForAdvancedZone(DataCenter dc) {
+        if (dc.getNetworkType() != NetworkType.Advanced) {
+            throw new CloudRuntimeException("Zone " + dc + " is not advanced.");
+        }
+
+        if (dc.isSecurityGroupEnabled()) {
+            List<NetworkVO> networks = _networkDao.listByZoneSecurityGroup(dc.getId());
+            if (CollectionUtils.isEmpty(networks)) {
+                throw new CloudRuntimeException("Can not found security enabled network in SG Zone " + dc);
+            }
+
+            return networks.get(0);
+        }
+        else {
+            TrafficType defaultTrafficType = TrafficType.Public;
+            List<NetworkVO> defaultNetworks = _networkDao.listByZoneAndTrafficType(dc.getId(), defaultTrafficType);
+            // api should never allow this situation to happen
+            if (defaultNetworks.size() != 1) {
+                throw new CloudRuntimeException("Found " + defaultNetworks.size() + " networks of type " + defaultTrafficType + " when expect to find 1");
+            }
+
+            return defaultNetworks.get(0);
+        }
+    }
+
+    /**
+     * Get default network for secondary storage VM for starting up in a basic zone. Basic zones select
+     * the Guest network whether or not the zone is SG-enabled.
+     * @param dc - The zone.
+     * @return The default network according to the zone's network selection rules.
+     * @throws CloudRuntimeException - If the zone is not a valid choice or a network couldn't be found.
+     */
+    protected NetworkVO getDefaultNetworkForBasicZone(DataCenter dc) {
+        if (dc.getNetworkType() != NetworkType.Basic) {
+            throw new CloudRuntimeException("Zone " + dc + "is not basic.");
+        }
+
+        TrafficType defaultTrafficType = TrafficType.Guest;
+        List<NetworkVO> defaultNetworks = _networkDao.listByZoneAndTrafficType(dc.getId(), defaultTrafficType);
+        // api should never allow this situation to happen
+        if (defaultNetworks.size() != 1) {
+            throw new CloudRuntimeException("Found " + defaultNetworks.size() + " networks of type " + defaultTrafficType + " when expect to find 1");
+        }
+
+        return defaultNetworks.get(0);
+    }
+
     protected Map<String, Object> createSecStorageVmInstance(long dataCenterId, SecondaryStorageVm.Role role) {
         DataStore secStore = _dataStoreMgr.getImageStore(dataCenterId);
         if (secStore == null) {
@@ -536,26 +607,7 @@ public class SecondaryStorageManagerImpl extends ManagerBase implements Secondar
         DataCenterDeployment plan = new DataCenterDeployment(dataCenterId);
         DataCenter dc = _dcDao.findById(plan.getDataCenterId());
 
-        NetworkVO defaultNetwork = null;
-        if (dc.getNetworkType() == NetworkType.Advanced && dc.isSecurityGroupEnabled()) {
-            List<NetworkVO> networks = _networkDao.listByZoneSecurityGroup(dataCenterId);
-            if (networks == null || networks.size() == 0) {
-                throw new CloudRuntimeException("Can not found security enabled network in SG Zone " + dc);
-            }
-            defaultNetwork = networks.get(0);
-        } else {
-            TrafficType defaultTrafficType = TrafficType.Public;
-
-            if (dc.getNetworkType() == NetworkType.Basic || dc.isSecurityGroupEnabled()) {
-                defaultTrafficType = TrafficType.Guest;
-            }
-            List<NetworkVO> defaultNetworks = _networkDao.listByZoneAndTrafficType(dataCenterId, defaultTrafficType);
-            // api should never allow this situation to happen
-            if (defaultNetworks.size() != 1) {
-                throw new CloudRuntimeException("Found " + defaultNetworks.size() + " networks of type " + defaultTrafficType + " when expect to find 1");
-            }
-            defaultNetwork = defaultNetworks.get(0);
-        }
+        NetworkVO defaultNetwork = getDefaultNetworkForCreation(dc);
 
         List<? extends NetworkOffering> offerings = null;
         if (_sNwMgr.isStorageIpRangeAvailable(dataCenterId)) {

services/secondary-storage/controller/test/org/apache/cloudstack/secondarystorage/SecondaryStorageManagerTest.java

@@ -0,0 +1,192 @@
+// 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 org.apache.cloudstack.secondarystorage;
+
+import java.util.Collections;
+
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.InjectMocks;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.MockitoAnnotations;
+
+import com.cloud.dc.DataCenterVO;
+import com.cloud.dc.DataCenter.NetworkType;
+import com.cloud.dc.dao.DataCenterDao;
+import com.cloud.network.Networks.TrafficType;
+import com.cloud.network.dao.NetworkDao;
+import com.cloud.network.dao.NetworkVO;
+import com.cloud.utils.exception.CloudRuntimeException;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.anyLong;
+import static org.mockito.Mockito.when;
+import static org.mockito.AdditionalMatchers.not;
+import static org.mockito.Mockito.eq;
+
+public class SecondaryStorageManagerTest {
+    @Mock
+    DataCenterDao _dcDao;
+
+    @Mock
+    NetworkDao _networkDao;
+
+    @InjectMocks
+    SecondaryStorageManagerImpl _ssMgr = new SecondaryStorageManagerImpl();
+
+    @Before
+    public void initMocks() {
+        MockitoAnnotations.initMocks(this);
+    }
+
+    @Test
+    public void getDefaultNetworkForAdvancedNonSG() {
+        DataCenterVO dc = mock(DataCenterVO.class);
+        when(dc.getNetworkType()).thenReturn(NetworkType.Advanced);
+        when(dc.isSecurityGroupEnabled()).thenReturn(false);
+
+        when(_dcDao.findById(Mockito.anyLong())).thenReturn(dc);
+
+        NetworkVO network = Mockito.mock(NetworkVO.class);
+        NetworkVO badNetwork = Mockito.mock(NetworkVO.class);
+        when(_networkDao.listByZoneAndTrafficType(anyLong(), eq(TrafficType.Public)))
+                    .thenReturn(Collections.singletonList(network));
+
+        when(_networkDao.listByZoneAndTrafficType(anyLong(), not(eq(TrafficType.Public))))
+        .thenReturn(Collections.singletonList(badNetwork));
+
+        when(_networkDao.listByZoneSecurityGroup(anyLong()))
+                    .thenReturn(Collections.singletonList(badNetwork));
+
+        NetworkVO returnedNetwork = _ssMgr.getDefaultNetworkForAdvancedZone(dc);
+
+        Assert.assertNotNull(returnedNetwork);
+        Assert.assertEquals(network, returnedNetwork);
+        Assert.assertNotEquals(badNetwork, returnedNetwork);
+    }
+
+    @Test
+    public void getDefaultNetworkForAdvancedSG() {
+        DataCenterVO dc = Mockito.mock(DataCenterVO.class);
+        when(dc.getNetworkType()).thenReturn(NetworkType.Advanced);
+        when(dc.isSecurityGroupEnabled()).thenReturn(true);
+
+        when(_dcDao.findById(Mockito.anyLong())).thenReturn(dc);
+
+        NetworkVO network = Mockito.mock(NetworkVO.class);
+        NetworkVO badNetwork = Mockito.mock(NetworkVO.class);
+        when(_networkDao.listByZoneAndTrafficType(anyLong(), any(TrafficType.class)))
+                    .thenReturn(Collections.singletonList(badNetwork));
+
+        when(_networkDao.listByZoneSecurityGroup(anyLong()))
+                    .thenReturn(Collections.singletonList(network));
+
+        NetworkVO returnedNetwork = _ssMgr.getDefaultNetworkForAdvancedZone(dc);
+
+        Assert.assertEquals(network, returnedNetwork);
+        Assert.assertNotEquals(badNetwork, returnedNetwork);
+    }
+
+    @Test
+    public void getDefaultNetworkForBasicNonSG() {
+        DataCenterVO dc = Mockito.mock(DataCenterVO.class);
+        when(dc.getNetworkType()).thenReturn(NetworkType.Basic);
+        when(dc.isSecurityGroupEnabled()).thenReturn(false);
+
+        when(_dcDao.findById(Mockito.anyLong())).thenReturn(dc);
+
+        NetworkVO network = Mockito.mock(NetworkVO.class);
+        NetworkVO badNetwork = Mockito.mock(NetworkVO.class);
+        when(_networkDao.listByZoneAndTrafficType(anyLong(), eq(TrafficType.Guest)))
+                    .thenReturn(Collections.singletonList(network));
+
+        when(_networkDao.listByZoneAndTrafficType(anyLong(), not(eq(TrafficType.Guest))))
+                    .thenReturn(Collections.singletonList(badNetwork));
+
+        NetworkVO returnedNetwork = _ssMgr.getDefaultNetworkForBasicZone(dc);
+
+        Assert.assertNotNull(returnedNetwork);
+        Assert.assertEquals(network, returnedNetwork);
+        Assert.assertNotEquals(badNetwork, returnedNetwork);
+    }
+
+    @Test
+    public void getDefaultNetworkForBasicSG() {
+        DataCenterVO dc = Mockito.mock(DataCenterVO.class);
+        when(dc.getNetworkType()).thenReturn(NetworkType.Basic);
+        when(dc.isSecurityGroupEnabled()).thenReturn(true);
+
+        when(_dcDao.findById(Mockito.anyLong())).thenReturn(dc);
+
+        NetworkVO network = Mockito.mock(NetworkVO.class);
+        NetworkVO badNetwork = Mockito.mock(NetworkVO.class);
+        when(_networkDao.listByZoneAndTrafficType(anyLong(), eq(TrafficType.Guest)))
+                    .thenReturn(Collections.singletonList(network));
+
+        when(_networkDao.listByZoneAndTrafficType(anyLong(), not(eq(TrafficType.Guest))))
+                    .thenReturn(Collections.singletonList(badNetwork));
+
+        NetworkVO returnedNetwork = _ssMgr.getDefaultNetworkForBasicZone(dc);
+
+        Assert.assertNotNull(returnedNetwork);
+        Assert.assertEquals(network, returnedNetwork);
+        Assert.assertNotEquals(badNetwork, returnedNetwork);
+    }
+
+    //also test invalid input
+    @Test(expected=CloudRuntimeException.class)
+    public void getDefaultNetworkForBasicSGWrongZoneType() {
+        DataCenterVO dc = Mockito.mock(DataCenterVO.class);
+        when(dc.getNetworkType()).thenReturn(NetworkType.Advanced);
+        when(dc.isSecurityGroupEnabled()).thenReturn(true);
+
+        when(_dcDao.findById(Mockito.anyLong())).thenReturn(dc);
+
+        NetworkVO network = Mockito.mock(NetworkVO.class);
+        NetworkVO badNetwork = Mockito.mock(NetworkVO.class);
+        when(_networkDao.listByZoneAndTrafficType(anyLong(), eq(TrafficType.Guest)))
+                    .thenReturn(Collections.singletonList(network));
+
+        when(_networkDao.listByZoneAndTrafficType(anyLong(), not(eq(TrafficType.Guest))))
+                    .thenReturn(Collections.singletonList(badNetwork));
+
+        _ssMgr.getDefaultNetworkForBasicZone(dc);
+    }
+
+    @Test(expected=CloudRuntimeException.class)
+    public void getDefaultNetworkForAdvancedWrongZoneType() {
+        DataCenterVO dc = Mockito.mock(DataCenterVO.class);
+        when(dc.getNetworkType()).thenReturn(NetworkType.Basic);
+        when(dc.isSecurityGroupEnabled()).thenReturn(true);
+
+        when(_dcDao.findById(Mockito.anyLong())).thenReturn(dc);
+
+        NetworkVO network = Mockito.mock(NetworkVO.class);
+        NetworkVO badNetwork = Mockito.mock(NetworkVO.class);
+        when(_networkDao.listByZoneAndTrafficType(anyLong(), any(TrafficType.class)))
+                    .thenReturn(Collections.singletonList(badNetwork));
+
+        when(_networkDao.listByZoneSecurityGroup(anyLong()))
+                    .thenReturn(Collections.singletonList(network));
+
+        _ssMgr.getDefaultNetworkForAdvancedZone(dc);
+    }
+}