Changing DHCP allocation map key to account for VLAN-id.

This way different services for the same subscriber, which have the same macAddress
but different vlanIds, won't overwrite each other. Allocations will also show the vlanId.
Also allocations were previously done only for DHCPACK for messages from the server.
Now updating allocations on all messages from dhcp-server.
Finally, updating unit tests to check dhcp-allocations, and cleaning up the logs by ignoring port-stats.

Change-Id: Ib48039ddcea0de7d3a60a7a6c43df40b932c4811
diff --git a/api/src/main/java/org/opencord/dhcpl2relay/DhcpAllocationInfo.java b/api/src/main/java/org/opencord/dhcpl2relay/DhcpAllocationInfo.java
index 770b2f0..6b5e3cd 100755
--- a/api/src/main/java/org/opencord/dhcpl2relay/DhcpAllocationInfo.java
+++ b/api/src/main/java/org/opencord/dhcpl2relay/DhcpAllocationInfo.java
@@ -15,13 +15,14 @@
  */
 package org.opencord.dhcpl2relay;
 
+import java.time.Instant;
+
 import org.onlab.packet.DHCP;
 import org.onlab.packet.IpAddress;
 import org.onlab.packet.MacAddress;
+import org.onlab.packet.VlanId;
 import org.onosproject.net.ConnectPoint;
 
-import java.time.Instant;
-
 /**
  * Information about DHCP Allocations.
  */
@@ -34,6 +35,7 @@
     private Instant allocationTime;
     private DHCP.MsgType type;
     private String subscriberId;
+    private VlanId vlanId;
 
     /**
      * Creates a new DHCP allocation info record.
@@ -43,16 +45,20 @@
      * @param circuitId option 82 information
      * @param macAddress MAC address of client
      * @param ip IP of client if allocated
+     * @param vlanId VLAN id of subscriber or service
      * @param subscriberId Sadis ID of a subscriber
      */
     public DhcpAllocationInfo(ConnectPoint location, DHCP.MsgType type,
-                              String circuitId, MacAddress macAddress, IpAddress ip, String subscriberId) {
+                              String circuitId, MacAddress macAddress,
+                              IpAddress ip, VlanId vlanId,
+                              String subscriberId) {
         this.location = location;
         this.type = type;
         this.circuitId = circuitId;
         this.macAddress = macAddress;
         this.ip = ip;
         this.allocationTime = Instant.now();
+        this.vlanId = vlanId;
         this.subscriberId = subscriberId;
     }
 
@@ -102,6 +108,15 @@
     }
 
     /**
+     * VLAN id of client or service.
+     *
+     * @return vlan id
+     */
+    public VlanId vlanId() {
+        return vlanId;
+    }
+
+    /**
      * SubscriberId of client if it has one.
      *
      * @return SubscriberId
@@ -118,4 +133,18 @@
     public Instant allocationTime() {
         return allocationTime;
     }
+
+    @Override
+    public String toString() {
+        return com.google.common.base.MoreObjects.toStringHelper(this)
+                .add("subscriberId", subscriberId)
+                .add("location", location)
+                .add("state", type)
+                .add("macAddress", macAddress)
+                .add("vlanId", vlanId)
+                .add("circuitId", circuitId)
+                .add("allocatedIP", ip)
+                .add("timestamp", allocationTime)
+                .toString();
+    }
 }
diff --git a/app/src/main/java/org/opencord/dhcpl2relay/impl/DhcpL2Relay.java b/app/src/main/java/org/opencord/dhcpl2relay/impl/DhcpL2Relay.java
index f32b675..1fc408d 100755
--- a/app/src/main/java/org/opencord/dhcpl2relay/impl/DhcpL2Relay.java
+++ b/app/src/main/java/org/opencord/dhcpl2relay/impl/DhcpL2Relay.java
@@ -15,10 +15,31 @@
  */
 package org.opencord.dhcpl2relay.impl;
 
-import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.Lists;
-import com.google.common.collect.Sets;
+import static org.onlab.packet.DHCP.DHCPOptionCode.OptionCode_MessageType;
+import static org.onlab.packet.MacAddress.valueOf;
+import static org.onosproject.net.config.basics.SubjectFactories.APP_SUBJECT_FACTORY;
+import static org.opencord.dhcpl2relay.impl.OsgiPropertyConstants.ENABLE_DHCP_BROADCAST_REPLIES;
+import static org.opencord.dhcpl2relay.impl.OsgiPropertyConstants.ENABLE_DHCP_BROADCAST_REPLIES_DEFAULT;
+import static org.opencord.dhcpl2relay.impl.OsgiPropertyConstants.OPTION_82;
+import static org.opencord.dhcpl2relay.impl.OsgiPropertyConstants.OPTION_82_DEFAULT;
+
+import java.io.ByteArrayOutputStream;
+import java.nio.ByteBuffer;
+import java.time.Instant;
+import java.util.ArrayList;
+import java.util.Dictionary;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Predicate;
+import java.util.stream.Collectors;
+
 import org.apache.commons.io.HexDump;
 import org.onlab.packet.DHCP;
 import org.onlab.packet.Ethernet;
@@ -92,30 +113,10 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.io.ByteArrayOutputStream;
-import java.nio.ByteBuffer;
-import java.time.Instant;
-import java.util.ArrayList;
-import java.util.Dictionary;
-import java.util.List;
-import java.util.Map;
-import java.util.Optional;
-import java.util.Set;
-import java.util.UUID;
-import java.util.concurrent.Executors;
-import java.util.concurrent.ScheduledExecutorService;
-import java.util.concurrent.ScheduledFuture;
-import java.util.concurrent.atomic.AtomicReference;
-import java.util.function.Predicate;
-import java.util.stream.Collectors;
-
-import static org.onlab.packet.DHCP.DHCPOptionCode.OptionCode_MessageType;
-import static org.onlab.packet.MacAddress.valueOf;
-import static org.onosproject.net.config.basics.SubjectFactories.APP_SUBJECT_FACTORY;
-import static org.opencord.dhcpl2relay.impl.OsgiPropertyConstants.ENABLE_DHCP_BROADCAST_REPLIES;
-import static org.opencord.dhcpl2relay.impl.OsgiPropertyConstants.ENABLE_DHCP_BROADCAST_REPLIES_DEFAULT;
-import static org.opencord.dhcpl2relay.impl.OsgiPropertyConstants.OPTION_82;
-import static org.opencord.dhcpl2relay.impl.OsgiPropertyConstants.OPTION_82_DEFAULT;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Lists;
+import com.google.common.collect.Sets;
 
 /**
  * DHCP Relay Agent Application Component.
@@ -698,7 +699,7 @@
             return null;
         }
 
-        //process the dhcp packet before sending to server
+        // process the dhcp packet before relaying to server or client
         private void processDhcpPacket(PacketContext context, Ethernet packet,
                                        DHCP dhcpPayload) {
             if (dhcpPayload == null) {
@@ -814,6 +815,7 @@
             IPv4 ipv4Packet = (IPv4) etherReply.getPayload();
             UDP udpPacket = (UDP) ipv4Packet.getPayload();
             DHCP dhcpPacket = (DHCP) udpPacket.getPayload();
+            ConnectPoint inPort = context.inPacket().receivedFrom();
 
             if (enableDhcpBroadcastReplies) {
                 // We want the reply to come back as a L2 broadcast
@@ -821,6 +823,7 @@
             }
 
             MacAddress clientMac = MacAddress.valueOf(dhcpPacket.getClientHardwareAddress());
+            VlanId clientVlan = VlanId.vlanId(ethernetPacket.getVlanID());
             IpAddress clientIp = IpAddress.valueOf(dhcpPacket.getClientIPAddress());
 
             SubscriberAndDeviceInformation entry = getSubscriber(context);
@@ -832,19 +835,20 @@
             UniTagInformation uniTagInformation = getUnitagInformationFromPacketContext(context, entry);
             if (uniTagInformation == null) {
                 log.warn("Missing service information for connectPoint {} / cTag {}",
-                         context.inPacket().receivedFrom(), context.inPacket().parsed().getVlanID());
+                         inPort, clientVlan);
                 return null;
             }
 
             DhcpAllocationInfo info = new DhcpAllocationInfo(
-                    context.inPacket().receivedFrom(), dhcpPacket.getPacketType(),
-                    entry.circuitId(), clientMac, clientIp, entry.id());
+                    inPort, dhcpPacket.getPacketType(), entry.circuitId(),
+                    clientMac, clientIp, clientVlan, entry.id());
 
-            String key = getUniqueUuidFromString(entry.id() + info.macAddress());
+            String key = getUniqueUuidFromString(entry.id() + clientMac
+                                                 + clientVlan);
             allocations.put(key, info);
 
             post(new DhcpL2RelayEvent(DhcpL2RelayEvent.Type.UPDATED, info,
-                                      context.inPacket().receivedFrom()));
+                                      inPort));
             if (option82) {
                 DHCP dhcpPacketWithOption82 = addOption82(dhcpPacket, entry);
                 udpPacket.setPayload(dhcpPacketWithOption82);
@@ -881,31 +885,30 @@
             IPv4 ipv4Packet = (IPv4) etherReply.getPayload();
             UDP udpPacket = (UDP) ipv4Packet.getPayload();
             DHCP dhcpPayload = (DHCP) udpPacket.getPayload();
+            VlanId innerVlan = VlanId.vlanId(ethernetPacket.getVlanID());
 
             MacAddress dstMac = valueOf(dhcpPayload.getClientHardwareAddress());
             ConnectPoint subsCp = getConnectPointOfClient(dstMac, context);
             // If we can't find the subscriber, can't process further
             if (subsCp == null) {
-                log.warn("Couldn't find connection point for mac address {} DHCPOFFERs won't be delivered", dstMac);
+                log.warn("Couldn't find subscriber, service or host info for mac"
+                        + " address {} .. DHCP packet won't be delivered", dstMac);
                 return null;
             }
 
             SubscriberAndDeviceInformation entry = getSubscriberInfoFromServer(dhcpPayload, context);
-
-            // if it's an ACK packet store the information for display purpose
-            if ((getDhcpPacketType(dhcpPayload) == DHCP.MsgType.DHCPACK) && (entry != null)) {
-
+            if (entry != null) {
                 IpAddress ip = IpAddress.valueOf(dhcpPayload.getYourIPAddress());
-                //storeDHCPAllocationInfo
+                // store DHCPAllocationInfo
                 DhcpAllocationInfo info = new DhcpAllocationInfo(subsCp,
-                     dhcpPayload.getPacketType(), entry.circuitId(), dstMac, ip, entry.id());
-
-                String key = getUniqueUuidFromString(entry.id() + info.macAddress());
+                     dhcpPayload.getPacketType(), entry.circuitId(), dstMac, ip,
+                     innerVlan, entry.id());
+                String key = getUniqueUuidFromString(entry.id()
+                                                     + info.macAddress() + innerVlan);
                 allocations.put(key, info);
 
                 post(new DhcpL2RelayEvent(DhcpL2RelayEvent.Type.UPDATED, info, subsCp));
-            } // end storing of info
-
+            }
 
             UniTagInformation uniTagInformation = getUnitagInformationFromPacketContext(context, entry);
             if (uniTagInformation == null) {
@@ -948,7 +951,8 @@
                 return null;
             }
             for (Host h : hosts) {
-                // if more than one,
+                // if more than one (for example, multiple services with same
+                // mac-address but different service VLANs (inner/C vlans)
                 // find the connect point which has an valid entry in SADIS
                 ConnectPoint cp = new ConnectPoint(h.location().deviceId(),
                                                    h.location().port());
@@ -959,20 +963,21 @@
                     log.warn("Subscriber info not found for {}", cp);
                     return null;
                 }
-
+                // check for cvlan in subscriber's uniTagInfo list
                 UniTagInformation uniTagInformation = getUnitagInformationFromPacketContext(context, sub);
                 if (uniTagInformation != null) {
                     return cp;
                 }
             }
             // no sadis config found for this connectPoint/vlan
-            log.warn("Missing service information for connectPoint {} / cTag {}",
+            log.warn("Missing service information for dhcp packet received from"
+                    + " {} with cTag {} .. cannot relay to client",
                      context.inPacket().receivedFrom(), context.inPacket().parsed().getVlanID());
 
             return null;
         }
 
-        //send the response to the requester host.
+        // send the response to the requester host (client)
         private void sendReply(Ethernet ethPacket, DHCP dhcpPayload, PacketContext context) {
             MacAddress descMac = valueOf(dhcpPayload.getClientHardwareAddress());
             ConnectPoint subCp = getConnectPointOfClient(descMac, context);
@@ -1068,11 +1073,13 @@
                 .forEach(allocations::remove);
     }
 
+    @Override
     public void clearAllocations() {
         allocations.clear();
     }
 
 
+    @Override
     public boolean removeAllocationsByConnectPoint(ConnectPoint cp) {
         boolean removed = false;
         for (String key : allocations.keySet()) {
@@ -1123,8 +1130,13 @@
             if (!isLocalLeader(deviceId)) {
                 return;
             }
+            // ignore stats
+            if (event.type().equals(DeviceEvent.Type.PORT_STATS_UPDATED)) {
+                return;
+            }
 
-            log.debug("Handling event {}", event);
+            log.debug("Device Event received for {} event {}", event.subject(),
+                      event.type());
 
             switch (event.type()) {
                 case DEVICE_REMOVED:
@@ -1152,11 +1164,6 @@
                 default:
                     break;
             }
-            if (log.isTraceEnabled() &&
-                    !event.type().equals(DeviceEvent.Type.PORT_STATS_UPDATED)) {
-                log.trace("Device Event received for {} event {}",
-                          event.subject(), event.type());
-            }
             if (!useOltUplink) {
                 if (dhcpServerConnectPoint.get() == null) {
                     switch (event.type()) {
diff --git a/app/src/test/java/org/opencord/dhcpl2relay/impl/DhcpL2RelayTest.java b/app/src/test/java/org/opencord/dhcpl2relay/impl/DhcpL2RelayTest.java
index 4106ef1..e37470b 100755
--- a/app/src/test/java/org/opencord/dhcpl2relay/impl/DhcpL2RelayTest.java
+++ b/app/src/test/java/org/opencord/dhcpl2relay/impl/DhcpL2RelayTest.java
@@ -15,7 +15,14 @@
  */
 package org.opencord.dhcpl2relay.impl;
 
-import com.google.common.collect.Lists;
+import static org.easymock.EasyMock.createMock;
+import static org.junit.Assert.assertEquals;
+import static org.slf4j.LoggerFactory.getLogger;
+
+import java.nio.ByteBuffer;
+import java.util.List;
+import java.util.Map;
+
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -23,7 +30,9 @@
 import org.onlab.packet.DHCP;
 import org.onlab.packet.Ethernet;
 import org.onlab.packet.IPv4;
+import org.onlab.packet.MacAddress;
 import org.onlab.packet.UDP;
+import org.onlab.packet.VlanId;
 import org.onlab.packet.dhcp.DhcpOption;
 import org.onosproject.cfg.ComponentConfigService;
 import org.onosproject.cluster.ClusterServiceAdapter;
@@ -32,21 +41,20 @@
 import org.onosproject.net.flowobjective.FlowObjectiveServiceAdapter;
 import org.onosproject.store.cluster.messaging.ClusterCommunicationServiceAdapter;
 import org.onosproject.store.service.TestStorageService;
+import org.opencord.dhcpl2relay.DhcpAllocationInfo;
 import org.opencord.dhcpl2relay.DhcpL2RelayEvent;
 import org.opencord.dhcpl2relay.DhcpL2RelayStoreDelegate;
 import org.opencord.dhcpl2relay.impl.packet.DhcpOption82;
+import org.slf4j.Logger;
 
-import java.nio.ByteBuffer;
-import java.util.List;
-import java.util.Map;
-
-import static org.easymock.EasyMock.createMock;
-import static org.junit.Assert.assertEquals;
+import com.google.common.collect.Lists;
 
 public class DhcpL2RelayTest extends DhcpL2RelayTestBase {
 
     private DhcpL2Relay dhcpL2Relay;
     private SimpleDhcpL2RelayCountersStore store;
+    private final Logger log = getLogger(getClass());
+    Map<String, DhcpAllocationInfo> allocs;
 
     ComponentConfigService mockConfigService =
             createMock(ComponentConfigService.class);
@@ -91,6 +99,76 @@
         dhcpL2Relay.deactivate();
     }
 
+    private void checkAllocation(DHCP.MsgType messageType) {
+        ConnectPoint clientCp = ConnectPoint.deviceConnectPoint(OLT_DEV_ID + "/"
+                + String.valueOf(CLIENT_PORT));
+        allocs = dhcpL2Relay.getAllocationInfo();
+        assert allocs.size() == 1;
+        allocs.forEach((k, v) -> {
+            log.info("Allocation {} : {}", k, v);
+            assertEquals(v.type(), messageType);
+            assertEquals(v.macAddress(), CLIENT_MAC);
+            assertEquals(v.location(), clientCp);
+        });
+    }
+
+    @Test
+    public void testMultipleAllocations() throws Exception {
+        dhcpL2Relay.clearAllocations();
+        // Trigger a discover from one RG on port 32
+        MacAddress mac32 = MacAddress.valueOf("b4:96:91:0c:4f:e4");
+        VlanId vlan32a = VlanId.vlanId((short) 801);
+        Ethernet discover32a = constructDhcpDiscoverPacket(
+                                  mac32, vlan32a, (short) 0);
+        ConnectPoint client32 = ConnectPoint.deviceConnectPoint(
+                                                "of:0000b86a974385f7/32");
+        sendPacket(discover32a, client32);
+        allocs = dhcpL2Relay.getAllocationInfo();
+        assert allocs.size() == 1;
+
+        //Trigger a discover from another RG on port 4112
+        MacAddress mac4112 = MacAddress.valueOf("b4:96:91:0c:4f:c9");
+        VlanId vlan4112 = VlanId.vlanId((short) 101);
+        Ethernet discover4112 = constructDhcpDiscoverPacket(
+                                                            mac4112, vlan4112,
+                                                            (short) 0);
+        ConnectPoint client4112 = ConnectPoint.deviceConnectPoint(
+                "of:0000b86a974385f7/4112");
+        sendPacket(discover4112, client4112);
+        allocs = dhcpL2Relay.getAllocationInfo();
+        assert allocs.size() == 2;
+
+        // Trigger a discover for another service with a different vlan
+        // from the same UNI port 32
+        VlanId vlan32b = VlanId.vlanId((short) 802);
+        Ethernet discover32b = constructDhcpDiscoverPacket(
+                                  mac32, vlan32b, (short) 0);
+        sendPacket(discover32b, client32);
+        allocs = dhcpL2Relay.getAllocationInfo();
+        assert allocs.size() == 3;
+
+        allocs.forEach((k, v) -> {
+            log.info("Allocation {} : {}", k, v);
+            assertEquals(v.type(), DHCP.MsgType.DHCPDISCOVER);
+            if (v.subscriberId().equals("ALPHe3d1cea3-1")) {
+                assertEquals(v.macAddress(), mac32);
+                assertEquals(v.location(), client32);
+                if (!(v.vlanId().equals(vlan32a) || v.vlanId().equals(vlan32b))) {
+                    assert false;
+                }
+            } else if (v.subscriberId().equals("ALPHe3d1ceb7-1")) {
+                assertEquals(v.macAddress(), mac4112);
+                assertEquals(v.location(), client4112);
+                assertEquals(v.vlanId(), vlan4112);
+            } else {
+                assert false;
+            }
+        });
+
+        dhcpL2Relay.clearAllocations();
+        assert dhcpL2Relay.getAllocationInfo().size() == 0;
+    }
+
     /**
      * Tests the DHCP relay app by sending DHCP discovery Packet.
      *
@@ -98,13 +176,16 @@
      */
     @Test
     public void testDhcpDiscover()  throws Exception {
-        //  (1) Sending DHCP discover packet
+        // Sending DHCP Discover packet
+        dhcpL2Relay.clearAllocations();
         Ethernet discoverPacket = constructDhcpDiscoverPacket(CLIENT_MAC);
-
-        sendPacket(discoverPacket, ConnectPoint.deviceConnectPoint(OLT_DEV_ID + "/" + 1));
+        ConnectPoint clientCp = ConnectPoint.deviceConnectPoint(OLT_DEV_ID + "/"
+                + String.valueOf(CLIENT_PORT));
+        sendPacket(discoverPacket, clientCp);
 
         Ethernet discoverRelayed = (Ethernet) getPacket();
         compareClientPackets(discoverPacket, discoverRelayed);
+        checkAllocation(DHCP.MsgType.DHCPDISCOVER);
     }
 
     /**
@@ -114,13 +195,15 @@
      */
     @Test
     public void testDhcpRequest()  throws Exception {
-        //  (1) Sending DHCP discover packet
+        // Sending DHCP Request packet
         Ethernet requestPacket = constructDhcpRequestPacket(CLIENT_MAC);
-
-        sendPacket(requestPacket, ConnectPoint.deviceConnectPoint(OLT_DEV_ID + "/" + 1));
+        ConnectPoint clientCp = ConnectPoint.deviceConnectPoint(OLT_DEV_ID + "/"
+                + String.valueOf(CLIENT_PORT));
+        sendPacket(requestPacket, clientCp);
 
         Ethernet requestRelayed = (Ethernet) getPacket();
         compareClientPackets(requestPacket, requestRelayed);
+        checkAllocation(DHCP.MsgType.DHCPREQUEST);
     }
 
     /**
@@ -130,14 +213,15 @@
      */
     @Test
     public void testDhcpOffer() {
-        //  (1) Sending DHCP discover packet
+        // Sending DHCP Offer packet
         Ethernet offerPacket = constructDhcpOfferPacket(SERVER_MAC,
                 CLIENT_MAC, DESTINATION_ADDRESS_IP, DHCP_CLIENT_IP_ADDRESS);
-
-        sendPacket(offerPacket, ConnectPoint.deviceConnectPoint(OLT_DEV_ID + "/" + 1));
+        sendPacket(offerPacket, ConnectPoint.deviceConnectPoint(OLT_DEV_ID + "/"
+                + String.valueOf(UPLINK_PORT)));
 
         Ethernet offerRelayed = (Ethernet) getPacket();
         compareServerPackets(offerPacket, offerRelayed);
+        checkAllocation(DHCP.MsgType.DHCPOFFER);
     }
 
     /**
@@ -151,10 +235,12 @@
         Ethernet ackPacket = constructDhcpAckPacket(SERVER_MAC,
                 CLIENT_MAC, DESTINATION_ADDRESS_IP, DHCP_CLIENT_IP_ADDRESS);
 
-        sendPacket(ackPacket, ConnectPoint.deviceConnectPoint(OLT_DEV_ID + "/" + 1));
+        sendPacket(ackPacket, ConnectPoint.deviceConnectPoint(OLT_DEV_ID + "/"
+                + String.valueOf(UPLINK_PORT)));
 
         Ethernet ackRelayed = (Ethernet) getPacket();
         compareServerPackets(ackPacket, ackRelayed);
+        checkAllocation(DHCP.MsgType.DHCPACK);
     }
 
     /**
@@ -172,6 +258,7 @@
 
         Ethernet nakRelayed = (Ethernet) getPacket();
         compareServerPackets(nakPacket, nakRelayed);
+        checkAllocation(DHCP.MsgType.DHCPNAK);
     }
 
     /**
@@ -188,6 +275,7 @@
 
         Ethernet declineRelayed = (Ethernet) getPacket();
         compareClientPackets(declinePacket, declineRelayed);
+        checkAllocation(DHCP.MsgType.DHCPDECLINE);
     }
 
     /**
diff --git a/app/src/test/java/org/opencord/dhcpl2relay/impl/DhcpL2RelayTestBase.java b/app/src/test/java/org/opencord/dhcpl2relay/impl/DhcpL2RelayTestBase.java
index 3d88b38..20ac9f5 100755
--- a/app/src/test/java/org/opencord/dhcpl2relay/impl/DhcpL2RelayTestBase.java
+++ b/app/src/test/java/org/opencord/dhcpl2relay/impl/DhcpL2RelayTestBase.java
@@ -32,19 +32,18 @@
 import java.util.concurrent.Callable;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
-import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.ScheduledExecutorService;
-import java.util.concurrent.TimeoutException;
+import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 
-import com.google.common.collect.ImmutableSet;
 import org.onlab.packet.BasePacket;
 import org.onlab.packet.ChassisId;
 import org.onlab.packet.DHCP;
 import org.onlab.packet.Ethernet;
 import org.onlab.packet.IPv4;
-import org.onlab.packet.IpAddress;
 import org.onlab.packet.Ip4Address;
+import org.onlab.packet.IpAddress;
 import org.onlab.packet.MacAddress;
 import org.onlab.packet.UDP;
 import org.onlab.packet.VlanId;
@@ -94,9 +93,11 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import com.google.common.collect.ImmutableSet;
+
 
 /**
- * Common methods for AAA app testing.
+ * Common methods for DHCP app testing.
  */
 public class DhcpL2RelayTestBase {
     private final Logger log = LoggerFactory.getLogger(getClass());
@@ -114,6 +115,7 @@
     static final String DESTINATION_ADDRESS_IP = "1.1.1.1";
     static final String DHCP_CLIENT_IP_ADDRESS = "2.2.2.2";
     static final int UPLINK_PORT = 5;
+    static final int CLIENT_PORT = 1;
 
     static final String EXPECTED_IP = "10.2.0.2";
     static final String OLT_DEV_ID = "of:00000000000000aa";
@@ -123,7 +125,8 @@
     static final MacAddress OLT_MAC_ADDRESS = MacAddress.valueOf("01:02:03:04:05:06");
 
     static final ConnectPoint SERVER_CONNECT_POINT =
-            ConnectPoint.deviceConnectPoint("of:00000000000000aa/5");
+            ConnectPoint.deviceConnectPoint("of:00000000000000aa/" +
+                                            String.valueOf(UPLINK_PORT));
 
     static final DefaultAnnotations DEVICE_ANNOTATIONS = DefaultAnnotations.builder()
             .set(AnnotationKeys.PROTOCOL, SCHEME_NAME.toUpperCase()).build();
@@ -162,21 +165,30 @@
         private final Device device1 = new DhcpL2RelayTestBase.MockDevice(providerId, DEVICE_ID_1, Device.Type.SWITCH,
                 "foo.inc", "0", "0", OLT_DEV_ID, new ChassisId(),
                 DEVICE_ANNOTATIONS);
-
+        private final Device otherDevice = new DhcpL2RelayTestBase.MockDevice(
+                                                providerId,
+                                                DeviceId.deviceId("of:0000b86a974385f7"),
+                                                Device.Type.SWITCH,
+                                                "foo.inc", "0", "0", "EC1838000853", new ChassisId(),
+                                                DEVICE_ANNOTATIONS);
         @Override
         public Device getDevice(DeviceId devId) {
-            return device1;
-
+            if (devId.equals(DEVICE_ID_1)) {
+                return device1;
+            } else {
+                return otherDevice;
+            }
         }
 
         @Override
         public Port getPort(ConnectPoint cp) {
-            return new DhcpL2RelayTestBase.MockPort();
+            return new DhcpL2RelayTestBase.MockPort(cp);
         }
 
         @Override
         public Port getPort(DeviceId deviceId, PortNumber portNumber) {
-            return new DhcpL2RelayTestBase.MockPort();
+            return new DhcpL2RelayTestBase.MockPort(new ConnectPoint(deviceId,
+                                                                     portNumber));
         }
 
         @Override
@@ -200,7 +212,8 @@
         @Override
         public Set<Host> getHostsByMac(MacAddress mac) {
 
-            HostLocation loc = new HostLocation(DEVICE_ID_1, PortNumber.portNumber(22), 0);
+            HostLocation loc = new HostLocation(DEVICE_ID_1, PortNumber
+                    .portNumber(CLIENT_PORT), 0);
 
             IpAddress ip = IpAddress.valueOf("10.100.200.10");
 
@@ -219,7 +232,11 @@
     }
 
     class  MockPort implements Port {
+        private ConnectPoint cp;
 
+        public MockPort(ConnectPoint cp) {
+            this.cp = cp;
+        }
         @Override
         public boolean isEnabled() {
             return true;
@@ -249,7 +266,13 @@
 
             @Override
             public String value(String val) {
-                return "PON 1/1";
+                if (cp.port().toLong() == 32) {
+                    return "ALPHe3d1cea3-1";
+                } else if (cp.port().toLong() == 4112) {
+                    return "ALPHe3d1ceb7-1";
+                } else {
+                    return "PON 1/1";
+                }
             }
             @Override
             public Set<String> keys() {
@@ -296,13 +319,30 @@
         DhcpL2RelayTestBase.MockSubscriberAndDeviceInformation device =
                 new DhcpL2RelayTestBase.MockSubscriberAndDeviceInformation(OLT_DEV_ID, VlanId.NONE, VlanId.NONE, null,
                         null, OLT_MAC_ADDRESS, Ip4Address.valueOf("10.10.10.10"), UPLINK_PORT);
+        DhcpL2RelayTestBase.MockSubscriberAndDeviceInformation otherDevice =
+                new DhcpL2RelayTestBase.MockSubscriberAndDeviceInformation(
+                        "EC1838000853", VlanId.NONE, VlanId.NONE, null,
+                        null, OLT_MAC_ADDRESS, Ip4Address.valueOf("10.10.10.10"), UPLINK_PORT);
+
         DhcpL2RelayTestBase.MockSubscriberAndDeviceInformation sub =
                 new DhcpL2RelayTestBase.MockSubscriberAndDeviceInformation(CLIENT_ID_1, CLIENT_C_TAG,
                         CLIENT_S_TAG, CLIENT_NAS_PORT_ID, CLIENT_CIRCUIT_ID, null, null, -1);
+        DhcpL2RelayTestBase.MockSubscriberAndDeviceInformation sub32 =
+                new DhcpL2RelayTestBase.MockSubscriberAndDeviceInformation("ALPHe3d1cea3-1", VlanId.vlanId((short) 801),
+                        VlanId.vlanId((short) 111), CLIENT_NAS_PORT_ID, CLIENT_CIRCUIT_ID, null, null, -1);
+        DhcpL2RelayTestBase.MockSubscriberAndDeviceInformation sub4112 =
+                new DhcpL2RelayTestBase.MockSubscriberAndDeviceInformation("ALPHe3d1ceb7-1", VlanId.vlanId((short) 101),
+                        VlanId.vlanId((short) 222), CLIENT_NAS_PORT_ID, CLIENT_CIRCUIT_ID, null, null, -1);
         @Override
         public SubscriberAndDeviceInformation get(String id) {
             if (id.equals(OLT_DEV_ID)) {
                 return device;
+            } else if (id.equals("EC1838000853")) {
+                return otherDevice;
+            } else if (id.equals("ALPHe3d1cea3-1")) {
+                return sub32;
+            } else if (id.equals("ALPHe3d1ceb7-1")) {
+                return sub4112;
             } else {
                 return  sub;
             }
@@ -340,6 +380,18 @@
                     .setIsDhcpRequired(true)
                     .build();
             uniTagInformationList.add(uniTagInformation);
+
+            if (id.equals("ALPHe3d1cea3-1")) {
+                // a second service on the same UNI
+                uniTagInformation = new UniTagInformation.Builder()
+                        .setPonCTag(VlanId.vlanId(((short) (cTag.toShort() + 1))))
+                        .setPonSTag(sTag)
+                        .setUsPonCTagPriority(CLIENT_C_PBIT)
+                        .setIsDhcpRequired(true)
+                        .build();
+                uniTagInformationList.add(uniTagInformation);
+            }
+
             this.setUniTagList(uniTagInformationList);
         }
     }
@@ -587,22 +639,18 @@
         packetProcessor.process(context);
     }
 
-    /**
-     * Constructs an Ethernet packet with IP/UDP/DHCP payload.
-     *
-     * @return Ethernet packet
-     */
-    private Ethernet construcEthernetPacket(MacAddress srcMac, MacAddress dstMac,
+    private Ethernet constructEthernetPacket(MacAddress srcMac, MacAddress dstMac,
                                                 String dstIp, byte dhcpReqRsp,
                                                 MacAddress clientHwAddress,
-                                                Ip4Address dhcpClientIpAddress) {
+                                                Ip4Address dhcpClientIpAddress,
+                                                VlanId clientVlan, short clientPbit) {
         // Ethernet Frame.
         Ethernet ethPkt = new Ethernet();
         ethPkt.setSourceMACAddress(srcMac);
         ethPkt.setDestinationMACAddress(dstMac);
         ethPkt.setEtherType(Ethernet.TYPE_IPV4);
-        ethPkt.setVlanID(CLIENT_C_TAG.toShort());
-        ethPkt.setPriorityCode((byte) CLIENT_C_PBIT);
+        ethPkt.setVlanID(clientVlan.toShort());
+        ethPkt.setPriorityCode((byte) clientPbit);
 
         if (DHCP.OPCODE_REPLY == dhcpReqRsp) {
             ethPkt.setQinQPriorityCode((byte) 3);
@@ -646,6 +694,23 @@
         ethPkt.setPayload(ipv4Reply);
 
         return ethPkt;
+
+    }
+
+    /**
+     * Constructs an Ethernet packet with IP/UDP/DHCP payload and client
+     * VLAN information.
+     *
+     * @return Ethernet packet
+     */
+    private Ethernet construcEthernetPacket(MacAddress srcMac, MacAddress dstMac,
+                                            String dstIp, byte dhcpReqRsp,
+                                            MacAddress clientHwAddress,
+                                            Ip4Address dhcpClientIpAddress) {
+        return constructEthernetPacket(srcMac, dstMac, dstIp, dhcpReqRsp,
+                                       clientHwAddress, dhcpClientIpAddress,
+                                       CLIENT_C_TAG, CLIENT_C_PBIT);
+
     }
 
     /**
@@ -654,9 +719,8 @@
      * @return Ethernet packet
      */
     Ethernet constructDhcpDiscoverPacket(MacAddress clientMac) {
-
         Ethernet pkt = construcEthernetPacket(clientMac, MacAddress.BROADCAST,
-                "255.255.255.255", DHCP.OPCODE_REQUEST, MacAddress.NONE,
+                "255.255.255.255", DHCP.OPCODE_REQUEST, clientMac,
                 Ip4Address.valueOf("0.0.0.0"));
 
         IPv4 ipv4Packet = (IPv4) pkt.getPayload();
@@ -669,14 +733,33 @@
     }
 
     /**
+     * Constructs DHCP Discover Packet with client VLAN information.
+     *
+     * @return Ethernet packet
+     */
+    Ethernet constructDhcpDiscoverPacket(MacAddress clientMac, VlanId clientVlan,
+                                         short clientPbit) {
+        Ethernet pkt = constructEthernetPacket(clientMac, MacAddress.BROADCAST,
+                "255.255.255.255", DHCP.OPCODE_REQUEST, clientMac,
+                Ip4Address.valueOf("0.0.0.0"), clientVlan, clientPbit);
+
+        IPv4 ipv4Packet = (IPv4) pkt.getPayload();
+        UDP udpPacket = (UDP) ipv4Packet.getPayload();
+        DHCP dhcpPacket = (DHCP) udpPacket.getPayload();
+
+        dhcpPacket.setOptions(constructDhcpOptions(DHCP.MsgType.DHCPDISCOVER));
+
+        return pkt;
+    }
+
+    /**
      * Constructs DHCP Request Packet.
      *
      * @return Ethernet packet
      */
     Ethernet constructDhcpRequestPacket(MacAddress clientMac) {
-
         Ethernet pkt = construcEthernetPacket(clientMac, MacAddress.BROADCAST,
-                "255.255.255.255", DHCP.OPCODE_REQUEST, MacAddress.NONE,
+                "255.255.255.255", DHCP.OPCODE_REQUEST, clientMac,
                 Ip4Address.valueOf("0.0.0.0"));
 
         IPv4 ipv4Packet = (IPv4) pkt.getPayload();
@@ -695,7 +778,6 @@
      */
     Ethernet constructDhcpOfferPacket(MacAddress servMac, MacAddress clientMac,
                                            String ipAddress, String dhcpClientIpAddress) {
-
         Ethernet pkt = construcEthernetPacket(servMac, clientMac, ipAddress, DHCP.OPCODE_REPLY,
                 clientMac, Ip4Address.valueOf(dhcpClientIpAddress));
 
@@ -715,7 +797,6 @@
      */
     Ethernet constructDhcpAckPacket(MacAddress servMac, MacAddress clientMac,
                                            String ipAddress, String dhcpClientIpAddress) {
-
         Ethernet pkt = construcEthernetPacket(servMac, clientMac, ipAddress, DHCP.OPCODE_REPLY,
                 clientMac, Ip4Address.valueOf(dhcpClientIpAddress));
 
@@ -756,7 +837,7 @@
     Ethernet constructDhcpDeclinePacket(MacAddress clientMac) {
 
         Ethernet pkt = construcEthernetPacket(clientMac, MacAddress.BROADCAST,
-                "255.255.255.255", DHCP.OPCODE_REQUEST, MacAddress.NONE,
+                "255.255.255.255", DHCP.OPCODE_REQUEST, clientMac,
                 Ip4Address.valueOf("0.0.0.0"));
 
         IPv4 ipv4Packet = (IPv4) pkt.getPayload();