Prepare the YANG models for RESTCONF
[lispflowmapping.git] / mappingservice / implementation / src / test / java / org / opendaylight / lispflowmapping / implementation / lisp / MapServerTest.java
index 10a4535d13c6cc8f141fa82ceaff3f3343d7bd66..775639180c8eb2cf2360949f9424f1b7dacb0a35 100644 (file)
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2013 Contextream, Inc. and others.  All rights reserved.
+ * Copyright (c) 2014 Contextream, Inc. and others.  All rights reserved.
  *
  * This program and the accompanying materials are made available under the
  * terms of the Eclipse Public License v1.0 which accompanies this distribution,
@@ -10,37 +10,40 @@ package org.opendaylight.lispflowmapping.implementation.lisp;
 
 import static org.junit.Assert.assertEquals;
 
-import java.util.HashMap;
-import java.util.Map;
+import java.util.ArrayList;
+import java.util.Arrays;
 
 import junitx.framework.ArrayAssert;
 
 import org.junit.Before;
 import org.junit.Test;
+import org.opendaylight.lispflowmapping.implementation.LispMappingService;
+import org.opendaylight.lispflowmapping.implementation.dao.MappingServiceKeyUtil;
 import org.opendaylight.lispflowmapping.implementation.serializer.MapRegisterSerializer;
+import org.opendaylight.lispflowmapping.implementation.util.LispAFIConvertor;
 import org.opendaylight.lispflowmapping.interfaces.dao.ILispDAO;
 import org.opendaylight.lispflowmapping.interfaces.dao.IMappingServiceKey;
 import org.opendaylight.lispflowmapping.interfaces.dao.MappingEntry;
-import org.opendaylight.lispflowmapping.interfaces.dao.MappingServiceKeyUtil;
-import org.opendaylight.lispflowmapping.interfaces.dao.MappingServiceValue;
+import org.opendaylight.lispflowmapping.interfaces.dao.MappingServiceRLOCGroup;
 import org.opendaylight.lispflowmapping.tools.junit.BaseTestCase;
-import org.opendaylight.lispflowmapping.type.lisp.EidToLocatorRecord;
-import org.opendaylight.lispflowmapping.type.lisp.LocatorRecord;
-import org.opendaylight.lispflowmapping.type.lisp.MapNotify;
-import org.opendaylight.lispflowmapping.type.lisp.MapRegister;
-import org.opendaylight.lispflowmapping.type.lisp.address.IMaskable;
-import org.opendaylight.lispflowmapping.type.lisp.address.LispAddress;
-import org.opendaylight.lispflowmapping.type.lisp.address.LispIpv4Address;
-import org.opendaylight.lispflowmapping.type.lisp.address.LispIpv6Address;
-import org.opendaylight.lispflowmapping.type.lisp.address.LispNoAddress;
+import org.opendaylight.yang.gen.v1.urn.opendaylight.lfm.control.plane.rev150314.EidToLocatorRecord.Action;
+import org.opendaylight.yang.gen.v1.urn.opendaylight.lfm.control.plane.rev150314.LispAFIAddress;
+import org.opendaylight.yang.gen.v1.urn.opendaylight.lfm.control.plane.rev150314.MapNotify;
+import org.opendaylight.yang.gen.v1.urn.opendaylight.lfm.control.plane.rev150314.MapRegister;
+import org.opendaylight.yang.gen.v1.urn.opendaylight.lfm.control.plane.rev150314.eidtolocatorrecords.EidToLocatorRecord;
+import org.opendaylight.yang.gen.v1.urn.opendaylight.lfm.control.plane.rev150314.eidtolocatorrecords.EidToLocatorRecordBuilder;
+import org.opendaylight.yang.gen.v1.urn.opendaylight.lfm.control.plane.rev150314.lispaddress.lispaddresscontainer.address.no.NoAddressBuilder;
+import org.opendaylight.yang.gen.v1.urn.opendaylight.lfm.control.plane.rev150314.locatorrecords.LocatorRecord;
+import org.opendaylight.yang.gen.v1.urn.opendaylight.lfm.control.plane.rev150314.locatorrecords.LocatorRecordBuilder;
+import org.opendaylight.yang.gen.v1.urn.opendaylight.lfm.control.plane.rev150314.mapregisternotification.MapRegisterBuilder;
 
 public class MapServerTest extends BaseTestCase {
 
-    private MapServer testedMapServer;
+    private LispMappingService testedMapServer;
     private ILispDAO lispDAO;
-    private MapRegister mapRegister;
-    private LispIpv4Address eid;
-    private LispIpv4Address rloc;
+    private MapRegisterBuilder mapRegisterBuilder;
+    private LispAFIAddress eid;
+    private LispAFIAddress rloc;
     private ValueSaverAction<MappingEntry<?>[]> mappingEntriesSaver;
     private MapRegister mapRegisterWithAuthentication;
 
@@ -49,15 +52,22 @@ public class MapServerTest extends BaseTestCase {
     public void before() throws Exception {
         super.before();
         lispDAO = context.mock(ILispDAO.class);
-        testedMapServer = new MapServer(lispDAO);
-        mapRegister = new MapRegister();
-        eid = new LispIpv4Address("10.31.0.5");
-        rloc = new LispIpv4Address(0xC0A8880A);
-        EidToLocatorRecord record = new EidToLocatorRecord();
-        record.setPrefix(eid).setMaskLength(32);
-        record.addLocator(new LocatorRecord().setLocator(rloc));
-        mapRegister.addEidToLocator(record);
-        mapRegisterWithAuthentication = MapRegisterSerializer.getInstance().deserialize(hexToByteBuffer("38 00 01 01 FF BB " //
+        testedMapServer = new LispMappingService();
+        testedMapServer.basicInit(lispDAO);
+        mapRegisterBuilder = new MapRegisterBuilder();
+        mapRegisterBuilder.setKeyId((short) 0);
+        mapRegisterBuilder.setAuthenticationData(new byte[0]);
+        eid = LispAFIConvertor.asIPAfiAddress("10.31.0.5");
+        rloc = LispAFIConvertor.asIPAfiAddress("192.168.136.10");
+        EidToLocatorRecordBuilder recordBuilder = new EidToLocatorRecordBuilder();
+        recordBuilder.setLispAddressContainer(LispAFIConvertor.toContainer(eid)).setMaskLength((short) 32);
+        recordBuilder.setLocatorRecord(new ArrayList<LocatorRecord>());
+        recordBuilder.getLocatorRecord().add(new LocatorRecordBuilder().setLispAddressContainer(LispAFIConvertor.toContainer(rloc)).build());
+        recordBuilder.setAction(Action.NoAction).setMapVersion((short) 0).setAuthoritative(false).setRecordTtl(60).setMaskLength((short) 32);
+        mapRegisterBuilder.setEidToLocatorRecord(new ArrayList<EidToLocatorRecord>());
+        mapRegisterBuilder.getEidToLocatorRecord().add(recordBuilder.build());
+        mapRegisterWithAuthentication = MapRegisterSerializer.getInstance().deserialize(hexToByteBuffer("38 00 01 01 FF BB "
+        //
                 + "00 00 00 00 00 00 00 01 00 14 2c 61 b9 c9 9a 20 " //
                 + "ba d8 f5 40 d3 55 6f 5f 6e 5a b2 0a bf b5 00 00 " //
                 + "00 0a 01 20 10 00 00 00 00 01 99 10 fe 01 01 64 " //
@@ -68,138 +78,254 @@ public class MapServerTest extends BaseTestCase {
 
     @Test
     public void handleMapRegister__NonSetMBit() throws Exception {
-        mapRegister.setWantMapNotify(false);
+        mapRegisterBuilder.setWantMapNotify(false);
 
-        addDAOExpectations(eid, 32);
-        assertNull(testedMapServer.handleMapRegister(mapRegister));
+        addDefaultPutAndGetExpectations(eid, 32);
+        assertNull(testedMapServer.handleMapRegister(mapRegisterBuilder.build(), false));
 
         MappingEntry<?>[] entries = mappingEntriesSaver.lastValue;
         assertEquals(1, entries.length);
 
-        assertEquals("value", entries[0].getKey());
-        assertEquals(rloc, ((MappingServiceValue) entries[0].getValue()).getRlocs().get(0).getRecord().getLocator());
+        assertEquals(AbstractLispComponent.ADDRESS_SUBKEY, entries[0].getKey());
+        assertEquals(1, ((MappingServiceRLOCGroup) entries[0].getValue()).getRecords().size());
+        assertEquals(rloc, LispAFIConvertor.toAFI(((MappingServiceRLOCGroup) entries[0].getValue()).getRecords().get(0).getLispAddressContainer()));
     }
 
     @Test
     public void handleMapRegisterIpv4__ValidNotifyEchoesRegister() throws Exception {
-        mapRegister.addEidToLocator(new EidToLocatorRecord().setPrefix(new LispNoAddress()));
-        mapRegister.setWantMapNotify(true);
-
-        addDAOExpectations(eid, 32);
-        MapNotify mapNotify = testedMapServer.handleMapRegister(mapRegister);
-        assertEquals(mapRegister.getEidToLocatorRecords(), mapNotify.getEidToLocatorRecords());
-        ArrayAssert.assertEquals(mapRegister.getAuthenticationData(), mapNotify.getAuthenticationData());
-        assertEquals(mapRegister.getKeyId(), mapNotify.getKeyId());
-        assertEquals(mapRegister.getNonce(), mapNotify.getNonce());
+        mapRegisterBuilder.getEidToLocatorRecord().add(
+                getDefaultEidToLocatorBuilder().setLispAddressContainer(LispAFIConvertor.toContainer(new NoAddressBuilder().build())).build());
+        mapRegisterBuilder.setWantMapNotify(true);
+
+        addDefaultPutAndGetExpectations(eid, 32);
+        addDefaultPutAndGetExpectations(new NoAddressBuilder().build(), 32);
+        MapRegister mr = mapRegisterBuilder.build();
+        MapNotify mapNotify = testedMapServer.handleMapRegister(mr, false);
+        assertEquals(mr.getEidToLocatorRecord(), mapNotify.getEidToLocatorRecord());
+        ArrayAssert.assertEquals(mr.getAuthenticationData(), mapNotify.getAuthenticationData());
+        assertEquals(mr.getKeyId(), mapNotify.getKeyId());
+        assertEquals(mr.getNonce(), mapNotify.getNonce());
     }
 
     @Test
     public void handleMapRegisterIpv4__CloneNotOwnYouClown() throws Exception {
-        mapRegister = new MapRegister();
-        mapRegister.setKeyId((byte) 0);
-        mapRegister.setWantMapNotify(true);
-        EidToLocatorRecord eidToLocator = new EidToLocatorRecord();
-        eid = new LispIpv4Address(1);
-        eidToLocator.setPrefix(eid);
-
-        LocatorRecord locator = new LocatorRecord();
-        locator.setLocator(new LispIpv4Address(2));
-        locator.setPriority((byte) 55);
-        eidToLocator.addLocator(locator);
+        mapRegisterBuilder = getDefaultMapRegisterBuilder();
+        mapRegisterBuilder.setWantMapNotify(true);
+        EidToLocatorRecordBuilder eidToLocatorBuilder = getDefaultEidToLocatorBuilder();
+        eidToLocatorBuilder.setMaskLength((short) 32);
+        eid = LispAFIConvertor.asIPAfiAddress("0.0.0.1");
+        eidToLocatorBuilder.setLispAddressContainer(LispAFIConvertor.toContainer(eid));
 
-        mapRegister.addEidToLocator(eidToLocator);
+        LocatorRecordBuilder locatorBuilder = getDefaultLocatorBuilder();
+        locatorBuilder.setLispAddressContainer(LispAFIConvertor.toContainer(LispAFIConvertor.asIPAfiAddress("0.0.0.2")));
+        locatorBuilder.setPriority((short) 55);
+        eidToLocatorBuilder.getLocatorRecord().add(locatorBuilder.build());
 
-        addDAOExpectations(eid, 0);
+        mapRegisterBuilder.getEidToLocatorRecord().add(eidToLocatorBuilder.build());
 
-        MapNotify mapNotify = testedMapServer.handleMapRegister(mapRegister);
+        addDefaultPutAndGetExpectations(eid, 32);
 
-        mapRegister.getEidToLocatorRecords().get(0).setPrefix(new LispIpv4Address(55));
-        mapRegister.getEidToLocatorRecords().get(0).getLocators().get(0).setPriority((byte) 1);
+        MapNotify mapNotify = testedMapServer.handleMapRegister(mapRegisterBuilder.build(), false);
 
-        EidToLocatorRecord actualEidToLocator = mapNotify.getEidToLocatorRecords().get(0);
-        assertEquals(new LispIpv4Address(1), actualEidToLocator.getPrefix());
-        assertEquals((byte) 55, actualEidToLocator.getLocators().get(0).getPriority());
+        EidToLocatorRecord actualEidToLocator = mapNotify.getEidToLocatorRecord().get(0);
+        assertEquals(LispAFIConvertor.toContainer(eid), actualEidToLocator.getLispAddressContainer());
+        assertEquals((byte) 55, actualEidToLocator.getLocatorRecord().get(0).getPriority().byteValue());
 
     }
 
     @Test
     public void handleMapRegisterIpv4__ValidMask() throws Exception {
         int mask = 16;
-        mapRegister = new MapRegister();
-        EidToLocatorRecord record = new EidToLocatorRecord();
-        record.setPrefix(eid).setMaskLength(mask);
-        record.addLocator(new LocatorRecord().setLocator(rloc));
-        mapRegister.addEidToLocator(record);
-        mapRegister.setWantMapNotify(true);
+        mapRegisterBuilder = getDefaultMapRegisterBuilder();
+        EidToLocatorRecordBuilder recordBuilder = getDefaultEidToLocatorBuilder();
+        recordBuilder.setLispAddressContainer(LispAFIConvertor.toContainer(eid)).setMaskLength((short) mask);
+        LocatorRecordBuilder locator = getDefaultLocatorBuilder();
+        locator.setLispAddressContainer(LispAFIConvertor.toContainer(rloc));
+        recordBuilder.getLocatorRecord().add(locator.build());
+        mapRegisterBuilder.getEidToLocatorRecord().add(recordBuilder.build());
+        mapRegisterBuilder.setWantMapNotify(true);
+        MapRegister mr = mapRegisterBuilder.build();
+
+        addDefaultPutAndGetExpectations(eid, mask);
+
+        MapNotify mapNotify = testedMapServer.handleMapRegister(mr, false);
+        assertEquals(LispAFIConvertor.toContainer(LispAFIConvertor.asIPAfiAddress("10.31.0.5")), mapNotify.getEidToLocatorRecord().get(0)
+                .getLispAddressContainer());
+        ArrayAssert.assertEquals(mr.getAuthenticationData(), mapNotify.getAuthenticationData());
+        assertEquals(mr.getEidToLocatorRecord(), mapNotify.getEidToLocatorRecord());
+        assertEquals(mr.getKeyId(), mapNotify.getKeyId());
+        assertEquals(mr.getNonce(), mapNotify.getNonce());
+    }
 
-        addDAOExpectations(eid, mask);
+    @Test
+    public void handleMapRegister__NonMaskable() throws Exception {
+        int mask = 16;
+        mapRegisterBuilder = getDefaultMapRegisterBuilder();
+        LispAFIAddress addr = LispAFIConvertor.asMacAfiAddress("01:01:01:01:01:01");
+        EidToLocatorRecordBuilder recordBuilder = getDefaultEidToLocatorBuilder();
+        LocatorRecordBuilder locatorBuilder = getDefaultLocatorBuilder();
+        locatorBuilder.setLispAddressContainer(LispAFIConvertor.toContainer(rloc));
+        recordBuilder.setLispAddressContainer(LispAFIConvertor.toContainer(addr)).setMaskLength((short) mask);
+        recordBuilder.getLocatorRecord().add(locatorBuilder.build());
+        mapRegisterBuilder.getEidToLocatorRecord().add(recordBuilder.build());
+        mapRegisterBuilder.setWantMapNotify(true);
+
+        addDefaultPutAndGetExpectations(addr, mask);
+
+        MapRegister mapRegister = mapRegisterBuilder.build();
+
+        MapNotify mapNotify = testedMapServer.handleMapRegister(mapRegister, false);
+        assertEquals(LispAFIConvertor.toContainer(addr), mapNotify.getEidToLocatorRecord().get(0).getLispAddressContainer());
+        ArrayAssert.assertEquals(mapRegister.getAuthenticationData(), mapNotify.getAuthenticationData());
+        assertEquals(mapRegister.getEidToLocatorRecord(), mapNotify.getEidToLocatorRecord());
+        assertEquals(mapRegister.getKeyId(), mapNotify.getKeyId());
+        assertEquals(mapRegister.getNonce(), mapNotify.getNonce());
+    }
 
-        MapNotify mapNotify = testedMapServer.handleMapRegister(mapRegister);
-        assertEquals(new LispIpv4Address("10.31.0.0"), mapNotify.getEidToLocatorRecords().get(0).getPrefix());
+    @Test
+    public void handleMapRegister__ZeroMask() throws Exception {
+        int mask = 0;
+        mapRegisterBuilder = getDefaultMapRegisterBuilder();
+        LispAFIAddress addr = LispAFIConvertor.asMacAfiAddress("01:01:01:01:01:01");
+        EidToLocatorRecordBuilder recordBuilder = getDefaultEidToLocatorBuilder();
+        LocatorRecordBuilder locatorBuilder = getDefaultLocatorBuilder();
+        locatorBuilder.setLispAddressContainer(LispAFIConvertor.toContainer(rloc));
+        recordBuilder.setLispAddressContainer(LispAFIConvertor.toContainer(addr)).setMaskLength((short) mask);
+        recordBuilder.getLocatorRecord().add(locatorBuilder.build());
+        mapRegisterBuilder.getEidToLocatorRecord().add(recordBuilder.build());
+        mapRegisterBuilder.setWantMapNotify(true);
+
+        addDefaultPutAndGetExpectations(addr, mask);
+
+        MapRegister mapRegister = mapRegisterBuilder.build();
+
+        MapNotify mapNotify = testedMapServer.handleMapRegister(mapRegister, false);
+        assertEquals(LispAFIConvertor.toContainer(addr), mapNotify.getEidToLocatorRecord().get(0).getLispAddressContainer());
         ArrayAssert.assertEquals(mapRegister.getAuthenticationData(), mapNotify.getAuthenticationData());
-        assertEquals(mapRegister.getEidToLocatorRecords(), mapNotify.getEidToLocatorRecords());
+        assertEquals(mapRegister.getEidToLocatorRecord(), mapNotify.getEidToLocatorRecord());
         assertEquals(mapRegister.getKeyId(), mapNotify.getKeyId());
         assertEquals(mapRegister.getNonce(), mapNotify.getNonce());
     }
 
+    @Test
+    public void handleMapRegisterIPv4__ZeroMask() throws Exception {
+        int mask = 0;
+        mapRegisterBuilder = getDefaultMapRegisterBuilder();
+        EidToLocatorRecordBuilder recordBuilder = getDefaultEidToLocatorBuilder();
+        recordBuilder.setLispAddressContainer(LispAFIConvertor.toContainer(eid)).setMaskLength((short) mask);
+        LocatorRecordBuilder locator = getDefaultLocatorBuilder();
+        locator.setLispAddressContainer(LispAFIConvertor.toContainer(rloc));
+        recordBuilder.getLocatorRecord().add(locator.build());
+        mapRegisterBuilder.getEidToLocatorRecord().add(recordBuilder.build());
+        mapRegisterBuilder.setWantMapNotify(true);
+        MapRegister mr = mapRegisterBuilder.build();
+
+        addDefaultPutAndGetExpectations(eid, mask);
+
+        MapNotify mapNotify = testedMapServer.handleMapRegister(mr, false);
+        assertEquals(LispAFIConvertor.toContainer(LispAFIConvertor.asIPAfiAddress("10.31.0.5")), mapNotify.getEidToLocatorRecord().get(0)
+                .getLispAddressContainer());
+        ArrayAssert.assertEquals(mr.getAuthenticationData(), mapNotify.getAuthenticationData());
+        assertEquals(mr.getEidToLocatorRecord(), mapNotify.getEidToLocatorRecord());
+        assertEquals(mr.getKeyId(), mapNotify.getKeyId());
+        assertEquals(mr.getNonce(), mapNotify.getNonce());
+    }
+
     @Test
     public void handleMapRegisterIpv4__ValidMask32() throws Exception {
         int mask = 32;
-        mapRegister = new MapRegister();
-        EidToLocatorRecord record = new EidToLocatorRecord();
-        record.setPrefix(new LispIpv4Address("10.31.0.5")).setMaskLength(mask);
-        record.addLocator(new LocatorRecord().setLocator(rloc));
-        mapRegister.addEidToLocator(record);
-        mapRegister.setWantMapNotify(true);
-
-        addDAOExpectations(eid, mask);
-
-        MapNotify mapNotify = testedMapServer.handleMapRegister(mapRegister);
-        assertEquals(new LispIpv4Address("10.31.0.5"), mapNotify.getEidToLocatorRecords().get(0).getPrefix());
-        assertEquals(mapRegister.getEidToLocatorRecords(), mapNotify.getEidToLocatorRecords());
+        mapRegisterBuilder = getDefaultMapRegisterBuilder();
+        EidToLocatorRecordBuilder recordBuilder = getDefaultEidToLocatorBuilder();
+        recordBuilder.setLispAddressContainer(LispAFIConvertor.toContainer(eid)).setMaskLength((short) mask);
+        LocatorRecordBuilder locator = getDefaultLocatorBuilder();
+        locator.setLispAddressContainer(LispAFIConvertor.toContainer(rloc));
+        recordBuilder.getLocatorRecord().add(locator.build());
+        mapRegisterBuilder.getEidToLocatorRecord().add(recordBuilder.build());
+        mapRegisterBuilder.setWantMapNotify(true);
+        MapRegister mr = mapRegisterBuilder.build();
+
+        addDefaultPutAndGetExpectations(eid, mask);
+
+        MapNotify mapNotify = testedMapServer.handleMapRegister(mr, false);
+        assertEquals(LispAFIConvertor.toContainer(LispAFIConvertor.asIPAfiAddress("10.31.0.5")), mapNotify.getEidToLocatorRecord().get(0)
+                .getLispAddressContainer());
+        ArrayAssert.assertEquals(mr.getAuthenticationData(), mapNotify.getAuthenticationData());
+        assertEquals(mr.getEidToLocatorRecord(), mapNotify.getEidToLocatorRecord());
+        assertEquals(mr.getKeyId(), mapNotify.getKeyId());
+        assertEquals(mr.getNonce(), mapNotify.getNonce());
+    }
+
+    @Test
+    public void handleMapRegisterIpv6__ValidMask96() throws Exception {
+        int mask = 96;
+        mapRegisterBuilder = getDefaultMapRegisterBuilder();
+        LispAFIAddress addr = LispAFIConvertor.asIPv6AfiAddress("1:1:1:1:1:1:1:0");
+        EidToLocatorRecordBuilder recordBuilder = getDefaultEidToLocatorBuilder();
+        LocatorRecordBuilder locatorBuilder = getDefaultLocatorBuilder();
+        locatorBuilder.setLispAddressContainer(LispAFIConvertor.toContainer(rloc));
+        recordBuilder.setLispAddressContainer(LispAFIConvertor.toContainer(addr)).setMaskLength((short) mask);
+        recordBuilder.getLocatorRecord().add(locatorBuilder.build());
+        mapRegisterBuilder.getEidToLocatorRecord().add(recordBuilder.build());
+        mapRegisterBuilder.setWantMapNotify(true);
+
+        addDefaultPutAndGetExpectations(addr, mask);
+
+        MapRegister mapRegister = mapRegisterBuilder.build();
+
+        MapNotify mapNotify = testedMapServer.handleMapRegister(mapRegister, false);
+        assertEquals(LispAFIConvertor.toContainer(addr), mapNotify.getEidToLocatorRecord().get(0).getLispAddressContainer());
         ArrayAssert.assertEquals(mapRegister.getAuthenticationData(), mapNotify.getAuthenticationData());
+        assertEquals(mapRegister.getEidToLocatorRecord(), mapNotify.getEidToLocatorRecord());
         assertEquals(mapRegister.getKeyId(), mapNotify.getKeyId());
         assertEquals(mapRegister.getNonce(), mapNotify.getNonce());
     }
 
     @Test
-    public void handleMapRegisterIpv6__ValidMask96() throws Exception {
-        int mask = 96;
-        mapRegister = new MapRegister();
-        EidToLocatorRecord record = new EidToLocatorRecord();
-        LispIpv6Address addr = new LispIpv6Address("1:1:1:1:1:1:1:0");
-        record.setPrefix(addr).setMaskLength(mask);
-        record.addLocator(new LocatorRecord().setLocator(rloc));
-        mapRegister.addEidToLocator(record);
-        mapRegister.setWantMapNotify(true);
-
-        addDAOExpectations(addr, mask);
-
-        MapNotify mapNotify = testedMapServer.handleMapRegister(mapRegister);
-        assertEquals(new LispIpv6Address("1:1:1:1:1:1:0:0"), mapNotify.getEidToLocatorRecords().get(0).getPrefix());
-        assertEquals(mapRegister.getEidToLocatorRecords(), mapNotify.getEidToLocatorRecords());
+    public void handleMapRegisterIpv6__ZeroMask() throws Exception {
+        int mask = 0;
+        mapRegisterBuilder = getDefaultMapRegisterBuilder();
+        LispAFIAddress addr = LispAFIConvertor.asIPv6AfiAddress("1:1:1:1:1:1:1:0");
+        EidToLocatorRecordBuilder recordBuilder = getDefaultEidToLocatorBuilder();
+        LocatorRecordBuilder locatorBuilder = getDefaultLocatorBuilder();
+        locatorBuilder.setLispAddressContainer(LispAFIConvertor.toContainer(rloc));
+        recordBuilder.setLispAddressContainer(LispAFIConvertor.toContainer(addr)).setMaskLength((short) mask);
+        recordBuilder.getLocatorRecord().add(locatorBuilder.build());
+        mapRegisterBuilder.getEidToLocatorRecord().add(recordBuilder.build());
+        mapRegisterBuilder.setWantMapNotify(true);
+
+        addDefaultPutAndGetExpectations(addr, mask);
+
+        MapRegister mapRegister = mapRegisterBuilder.build();
+
+        MapNotify mapNotify = testedMapServer.handleMapRegister(mapRegister, false);
+        assertEquals(LispAFIConvertor.toContainer(addr), mapNotify.getEidToLocatorRecord().get(0).getLispAddressContainer());
         ArrayAssert.assertEquals(mapRegister.getAuthenticationData(), mapNotify.getAuthenticationData());
+        assertEquals(mapRegister.getEidToLocatorRecord(), mapNotify.getEidToLocatorRecord());
         assertEquals(mapRegister.getKeyId(), mapNotify.getKeyId());
         assertEquals(mapRegister.getNonce(), mapNotify.getNonce());
     }
 
     @Test
-    public void handleMapRegisterIpv6__ValidMask32() throws Exception {
-        int mask = 32;
-        mapRegister = new MapRegister();
-        EidToLocatorRecord record = new EidToLocatorRecord();
-        LispIpv6Address addr = new LispIpv6Address("1:1:0:0:0:0:0:0");
-        record.setPrefix(addr).setMaskLength(mask);
-        record.addLocator(new LocatorRecord().setLocator(rloc));
-        mapRegister.addEidToLocator(record);
-        mapRegister.setWantMapNotify(true);
-
-        addDAOExpectations(addr, mask);
-
-        MapNotify mapNotify = testedMapServer.handleMapRegister(mapRegister);
-        assertEquals(new LispIpv6Address("1:1:0:0:0:0:0:0"), mapNotify.getEidToLocatorRecords().get(0).getPrefix());
-        assertEquals(mapRegister.getEidToLocatorRecords(), mapNotify.getEidToLocatorRecords());
+    public void handleMapRegisterIpv6__ValidMask48() throws Exception {
+        int mask = 48;
+        mapRegisterBuilder = getDefaultMapRegisterBuilder();
+        LispAFIAddress addr = LispAFIConvertor.asIPv6AfiAddress("1:1:1:1:1:1:1:0");
+        EidToLocatorRecordBuilder recordBuilder = getDefaultEidToLocatorBuilder();
+        LocatorRecordBuilder locatorBuilder = getDefaultLocatorBuilder();
+        locatorBuilder.setLispAddressContainer(LispAFIConvertor.toContainer(rloc));
+        recordBuilder.setLispAddressContainer(LispAFIConvertor.toContainer(addr)).setMaskLength((short) mask);
+        recordBuilder.getLocatorRecord().add(locatorBuilder.build());
+        mapRegisterBuilder.getEidToLocatorRecord().add(recordBuilder.build());
+        mapRegisterBuilder.setWantMapNotify(true);
+
+        addDefaultPutAndGetExpectations(addr, mask);
+
+        MapRegister mapRegister = mapRegisterBuilder.build();
+
+        MapNotify mapNotify = testedMapServer.handleMapRegister(mapRegister, false);
+        assertEquals(LispAFIConvertor.toContainer(addr), mapNotify.getEidToLocatorRecord().get(0).getLispAddressContainer());
         ArrayAssert.assertEquals(mapRegister.getAuthenticationData(), mapNotify.getAuthenticationData());
+        assertEquals(mapRegister.getEidToLocatorRecord(), mapNotify.getEidToLocatorRecord());
         assertEquals(mapRegister.getKeyId(), mapNotify.getKeyId());
         assertEquals(mapRegister.getNonce(), mapNotify.getNonce());
     }
@@ -207,57 +333,220 @@ public class MapServerTest extends BaseTestCase {
     @Test
     public void handleMapRegisterIpv6__ValidMask128() throws Exception {
         int mask = 128;
-        mapRegister = new MapRegister();
-        EidToLocatorRecord record = new EidToLocatorRecord();
-        LispIpv6Address addr = new LispIpv6Address("1:1:1:1:1:1:1:2");
-        record.setPrefix(addr).setMaskLength(mask);
-        record.addLocator(new LocatorRecord().setLocator(rloc));
-        mapRegister.addEidToLocator(record);
-        mapRegister.setWantMapNotify(true);
-
-        addDAOExpectations(addr, mask);
-
-        MapNotify mapNotify = testedMapServer.handleMapRegister(mapRegister);
-        assertEquals(new LispIpv6Address("1:1:1:1:1:1:1:2"), mapNotify.getEidToLocatorRecords().get(0).getPrefix());
-        assertEquals(mapRegister.getEidToLocatorRecords(), mapNotify.getEidToLocatorRecords());
+        mapRegisterBuilder = getDefaultMapRegisterBuilder();
+        LispAFIAddress addr = LispAFIConvertor.asIPv6AfiAddress("1:1:1:1:1:1:1:0");
+        EidToLocatorRecordBuilder recordBuilder = getDefaultEidToLocatorBuilder();
+        LocatorRecordBuilder locatorBuilder = getDefaultLocatorBuilder();
+        locatorBuilder.setLispAddressContainer(LispAFIConvertor.toContainer(rloc));
+        recordBuilder.setLispAddressContainer(LispAFIConvertor.toContainer(addr)).setMaskLength((short) mask);
+        recordBuilder.getLocatorRecord().add(locatorBuilder.build());
+        mapRegisterBuilder.getEidToLocatorRecord().add(recordBuilder.build());
+        mapRegisterBuilder.setWantMapNotify(true);
+
+        addDefaultPutAndGetExpectations(addr, mask);
+
+        MapRegister mapRegister = mapRegisterBuilder.build();
+
+        MapNotify mapNotify = testedMapServer.handleMapRegister(mapRegister, false);
+        assertEquals(LispAFIConvertor.toContainer(addr), mapNotify.getEidToLocatorRecord().get(0).getLispAddressContainer());
+        ArrayAssert.assertEquals(mapRegister.getAuthenticationData(), mapNotify.getAuthenticationData());
+        assertEquals(mapRegister.getEidToLocatorRecord(), mapNotify.getEidToLocatorRecord());
+        assertEquals(mapRegister.getKeyId(), mapNotify.getKeyId());
+        assertEquals(mapRegister.getNonce(), mapNotify.getNonce());
+    }
+
+    @Test
+    public void handleMapRegisterIPV4AndIpv6__ValidMask96() throws Exception {
+        int mask = 96;
+        mapRegisterBuilder = getDefaultMapRegisterBuilder();
+        EidToLocatorRecordBuilder recordBuilder0 = getDefaultEidToLocatorBuilder();
+        recordBuilder0.setLispAddressContainer(LispAFIConvertor.toContainer(eid)).setMaskLength((short) 32);
+        LispAFIAddress addr = LispAFIConvertor.asIPv6AfiAddress("1:1:1:1:1:1:1:0");
+        EidToLocatorRecordBuilder recordBuilder1 = getDefaultEidToLocatorBuilder();
+        LocatorRecordBuilder locatorBuilder = getDefaultLocatorBuilder();
+        locatorBuilder.setLispAddressContainer(LispAFIConvertor.toContainer(rloc));
+        locatorBuilder.setLispAddressContainer(LispAFIConvertor.toContainer(rloc));
+        recordBuilder1.setLispAddressContainer(LispAFIConvertor.toContainer(addr)).setMaskLength((short) mask);
+        recordBuilder1.getLocatorRecord().add(locatorBuilder.build());
+        mapRegisterBuilder.getEidToLocatorRecord().add(recordBuilder0.build());
+        mapRegisterBuilder.getEidToLocatorRecord().add(recordBuilder1.build());
+        mapRegisterBuilder.setWantMapNotify(true);
+
+        addDefaultPutAndGetExpectations(eid, 32);
+        addDefaultPutAndGetExpectations(addr, mask);
+
+        MapRegister mapRegister = mapRegisterBuilder.build();
+
+        MapNotify mapNotify = testedMapServer.handleMapRegister(mapRegister, false);
+        assertEquals(LispAFIConvertor.toContainer(eid), mapNotify.getEidToLocatorRecord().get(0).getLispAddressContainer());
+        assertEquals(LispAFIConvertor.toContainer(addr), mapNotify.getEidToLocatorRecord().get(1).getLispAddressContainer());
         ArrayAssert.assertEquals(mapRegister.getAuthenticationData(), mapNotify.getAuthenticationData());
+        assertEquals(mapRegister.getEidToLocatorRecord(), mapNotify.getEidToLocatorRecord());
         assertEquals(mapRegister.getKeyId(), mapNotify.getKeyId());
         assertEquals(mapRegister.getNonce(), mapNotify.getNonce());
     }
 
     @Test
     public void handleMapRegister__MultipleRLOCs() throws Exception {
-        addDAOExpectations(eid, 32);
+        addDefaultPutAndGetExpectations(eid, 32);
+
+        mapRegisterBuilder = getDefaultMapRegisterBuilder();
+        LispAFIAddress rloc0 = rloc;
+        LispAFIAddress rloc1 = LispAFIConvertor.asIPv6AfiAddress("0:0:0:0:0:0:0:7");
+        EidToLocatorRecordBuilder recordBuilder = getDefaultEidToLocatorBuilder();
+        recordBuilder.setLispAddressContainer(LispAFIConvertor.toContainer(eid));
+        LocatorRecordBuilder locatorBuilder1 = getDefaultLocatorBuilder();
+        locatorBuilder1.setLispAddressContainer(LispAFIConvertor.toContainer(rloc0));
+        LocatorRecordBuilder locatorBuilder2 = getDefaultLocatorBuilder();
+        locatorBuilder2.setLispAddressContainer(LispAFIConvertor.toContainer(rloc1));
+        recordBuilder.getLocatorRecord().add(locatorBuilder1.build());
+        recordBuilder.getLocatorRecord().add(locatorBuilder2.build());
+
+        mapRegisterBuilder.getEidToLocatorRecord().add(recordBuilder.build());
+
+        testedMapServer.handleMapRegister(mapRegisterBuilder.build(), false);
+
+        MappingEntry<?>[] entries = mappingEntriesSaver.lastValue;
+        assertEquals(1, entries.length);
 
-        LispIpv4Address rloc0 = rloc;
-        LispIpv6Address rloc1 = new LispIpv6Address("::7");
-        mapRegister.getEidToLocatorRecords().get(0).addLocator(new LocatorRecord().setLocator(rloc1));
+        assertEquals(AbstractLispComponent.ADDRESS_SUBKEY, entries[0].getKey());
+        assertEquals(LispAFIConvertor.toContainer(rloc0), ((MappingServiceRLOCGroup) entries[0].getValue()).getRecords().get(0)
+                .getLispAddressContainer());
+        assertEquals(LispAFIConvertor.toContainer(rloc1), ((MappingServiceRLOCGroup) entries[0].getValue()).getRecords().get(1)
+                .getLispAddressContainer());
 
-        testedMapServer.handleMapRegister(mapRegister);
+    }
+
+    @Test
+    public void handleMapRegister__MultipleTypes() throws Exception {
+        addDefaultPutAndGetExpectations(eid, 32);
+
+        mapRegisterBuilder = getDefaultMapRegisterBuilder();
+        LispAFIAddress rloc0 = rloc;
+        // LispAFIAddress rloc1 =
+        // LispAFIConvertor.asIPv6AfiAddress("0:0:0:0:0:0:0:7");
+        String subkey = "bla";
+        LispAFIAddress rloc1 = LispAFIConvertor
+                .asKeyValue(subkey, LispAFIConvertor.toPrimitive(LispAFIConvertor.asIPv6AfiAddress("0:0:0:0:0:0:0:7")));
+        EidToLocatorRecordBuilder recordBuilder = getDefaultEidToLocatorBuilder();
+        recordBuilder.setLispAddressContainer(LispAFIConvertor.toContainer(eid));
+        LocatorRecordBuilder locatorBuilder1 = getDefaultLocatorBuilder();
+        locatorBuilder1.setLispAddressContainer(LispAFIConvertor.toContainer(rloc0));
+        LocatorRecordBuilder locatorBuilder2 = getDefaultLocatorBuilder();
+        locatorBuilder2.setLispAddressContainer(LispAFIConvertor.toContainer(rloc1));
+        recordBuilder.getLocatorRecord().add(locatorBuilder1.build());
+        recordBuilder.getLocatorRecord().add(locatorBuilder2.build());
+
+        mapRegisterBuilder.getEidToLocatorRecord().add(recordBuilder.build());
+
+        testedMapServer.handleMapRegister(mapRegisterBuilder.build(), false);
+
+        MappingEntry<?>[] entries = mappingEntriesSaver.lastValue;
+        assertEquals(2, entries.length);
+
+        assertEquals(AbstractLispComponent.ADDRESS_SUBKEY, entries[0].getKey());
+        assertEquals(subkey, entries[1].getKey());
+        assertEquals(LispAFIConvertor.toContainer(rloc0), ((MappingServiceRLOCGroup) entries[0].getValue()).getRecords().get(0)
+                .getLispAddressContainer());
+        assertEquals(LispAFIConvertor.toContainer(rloc1), ((MappingServiceRLOCGroup) entries[1].getValue()).getRecords().get(0)
+                .getLispAddressContainer());
+
+    }
+
+    @Test
+    public void handleMapRegister__TestOverwrite() throws Exception {
+        addDefaultPutAndGetExpectations(eid, 32);
+
+        addEidToLocatorRecord();
 
         MappingEntry<?>[] entries = mappingEntriesSaver.lastValue;
         assertEquals(1, entries.length);
 
-        assertEquals("value", entries[0].getKey());
-        assertEquals(rloc0, ((MappingServiceValue) entries[0].getValue()).getRlocs().get(0).getRecord().getLocator());
-        assertEquals(rloc1, ((MappingServiceValue) entries[0].getValue()).getRlocs().get(1).getRecord().getLocator());
+        assertEquals(AbstractLispComponent.ADDRESS_SUBKEY, entries[0].getKey());
+        assertEquals(LispAFIConvertor.toContainer(rloc), ((MappingServiceRLOCGroup) entries[0].getValue()).getRecords().get(0)
+                .getLispAddressContainer());
+    }
+
+    @Test
+    public void handleMapRegister__TestDontOverwrite() throws Exception {
+        int hc = LispAFIConvertor.toContainer(rloc).getAddress().hashCode();
+
+        addDefaultPutAndGetExpectations(eid, 32);
+        testedMapServer.setOverwrite(false);
+
+        addEidToLocatorRecord();
+
+        MappingEntry<?>[] entries = mappingEntriesSaver.lastValue;
+        assertEquals(1, entries.length);
+
+        assertEquals(String.valueOf(hc), entries[0].getKey());
+        assertEquals(LispAFIConvertor.toContainer(rloc), ((MappingServiceRLOCGroup) entries[0].getValue()).getRecords().get(0)
+                .getLispAddressContainer());
+    }
+
+    private void addEidToLocatorRecord() {
+        EidToLocatorRecordBuilder recordBuilder = getDefaultEidToLocatorBuilder();
+        recordBuilder.setLispAddressContainer(LispAFIConvertor.toContainer(eid));
+        LocatorRecordBuilder locatorBuilder = getDefaultLocatorBuilder();
+        locatorBuilder.setLispAddressContainer(LispAFIConvertor.toContainer(rloc));
+        recordBuilder.getLocatorRecord().add(locatorBuilder.build());
+
+        mapRegisterBuilder = getDefaultMapRegisterBuilder();
+        mapRegisterBuilder.getEidToLocatorRecord().add(recordBuilder.build());
+
+        testedMapServer.handleMapRegister(mapRegisterBuilder.build(), false);
+    }
+
+    @Test
+    public void handleMapRegister__MultipleEIDs() throws Exception {
+        addDefaultPutAndGetExpectations(eid, 32);
+
+        LispAFIAddress rloc1 = LispAFIConvertor.asIPv6AfiAddress("0:0:0:0:0:0:0:7");
+        mapRegisterBuilder = getDefaultMapRegisterBuilder();
+        EidToLocatorRecordBuilder etlrBuilder = getDefaultEidToLocatorBuilder();
+        etlrBuilder.getLocatorRecord().add(getDefaultLocatorBuilder().setLispAddressContainer(LispAFIConvertor.toContainer(rloc1)).build());
+        mapRegisterBuilder.setWantMapNotify(true);
+        etlrBuilder.setLispAddressContainer(LispAFIConvertor.toContainer(eid));
+
+        EidToLocatorRecordBuilder etlr2Builder = getDefaultEidToLocatorBuilder();
+        LispAFIAddress address = LispAFIConvertor.asIPAfiAddress("1.1.1.1");
+        etlr2Builder.setLispAddressContainer(LispAFIConvertor.toContainer(address));
+        etlr2Builder.setMaskLength((short) 32);
+        int recordTtl = 5;
+        etlr2Builder.setRecordTtl(recordTtl);
+        etlr2Builder.getLocatorRecord().add(
+                getDefaultLocatorBuilder().setPriority((short) 10)
+                        .setLispAddressContainer(LispAFIConvertor.toContainer(LispAFIConvertor.asIPAfiAddress("2.2.2.2"))).build());
+        mapRegisterBuilder.getEidToLocatorRecord().add(etlrBuilder.build());
+        mapRegisterBuilder.getEidToLocatorRecord().add(etlr2Builder.build());
+        addDefaultPutAndGetExpectations(address, 32);
+        MapNotify mapNotify = testedMapServer.handleMapRegister(mapRegisterBuilder.build(), false);
+
+        assertEquals(LispAFIConvertor.toContainer(rloc1), mapNotify.getEidToLocatorRecord().get(0).getLocatorRecord().get(0)
+                .getLispAddressContainer());
+        assertEquals(LispAFIConvertor.toContainer(LispAFIConvertor.asIPAfiAddress("2.2.2.2")), mapNotify.getEidToLocatorRecord().get(1)
+                .getLocatorRecord().get(0).getLispAddressContainer());
+        assertEquals(LispAFIConvertor.toContainer(LispAFIConvertor.asIPAfiAddress("1.1.1.1")), mapNotify.getEidToLocatorRecord().get(1)
+                .getLispAddressContainer());
+        assertEquals(recordTtl, mapNotify.getEidToLocatorRecord().get(1).getRecordTtl().intValue());
 
     }
 
     @Test
     public void handleMapRegisterIpv4__ChekWrongPassword() throws Exception {
 
-        addDAOEpectations(new LispIpv4Address("153.16.254.1"), 32, 0, 31, "bla");
-        MapNotify mapNotify = testedMapServer.handleMapRegister(mapRegisterWithAuthentication);
+        addGetExpectations(LispAFIConvertor.asIPAfiAddress("153.16.254.1"), 32, 0, 31, "bla");
+        MapNotify mapNotify = testedMapServer.handleMapRegister(mapRegisterWithAuthentication, false);
         assertEquals(null, mapNotify);
     }
 
     @Test
     public void handleMapRegisterIpv4__ChcekNoPasswordAndThenPassword() throws Exception {
 
-        addDAOEpectations(new LispIpv4Address("153.16.254.1"), 32, 30, 25, "password");
-        MapNotify mapNotify = testedMapServer.handleMapRegister(mapRegisterWithAuthentication);
-        assertEquals(mapRegisterWithAuthentication.getEidToLocatorRecords(), mapNotify.getEidToLocatorRecords());
+        addGetExpectations(LispAFIConvertor.asIPAfiAddress("153.16.254.1"), 32, 30, 25, "password");
+        addPutExpectations(LispAFIConvertor.asIPAfiAddress("153.16.254.1"), 32);
+        MapNotify mapNotify = testedMapServer.handleMapRegister(mapRegisterWithAuthentication, false);
+        assertEquals(mapRegisterWithAuthentication.getEidToLocatorRecord(), mapNotify.getEidToLocatorRecord());
         assertEquals(mapRegisterWithAuthentication.getKeyId(), mapNotify.getKeyId());
         assertEquals(mapRegisterWithAuthentication.getNonce(), mapNotify.getNonce());
     }
@@ -265,38 +554,35 @@ public class MapServerTest extends BaseTestCase {
     @Test
     public void handleMapRegisterIpv4__ChcekNoPassword() throws Exception {
 
-        addDAOEpectations(new LispIpv4Address("153.16.254.1"), 32, 30, 0, "password");
-        MapNotify mapNotify = testedMapServer.handleMapRegister(mapRegisterWithAuthentication);
+        addGetExpectations(LispAFIConvertor.asIPAfiAddress("153.16.254.1"), 32, 30, 0, "password");
+        MapNotify mapNotify = testedMapServer.handleMapRegister(mapRegisterWithAuthentication, false);
         assertEquals(null, mapNotify);
     }
 
     @Test
     public void handleMapRegisterIpv4__ChcekNoreturn() throws Exception {
 
-        addDAOExpectations(new LispIpv4Address("153.16.254.1"), 32);
-        MapNotify mapNotify = testedMapServer.handleMapRegister(mapRegisterWithAuthentication);
+        addGetExpectations(LispAFIConvertor.asIPAfiAddress("153.16.254.1"), 32);
+        MapNotify mapNotify = testedMapServer.handleMapRegister(mapRegisterWithAuthentication, false);
         assertEquals(mapNotify, null);
     }
 
+    @SuppressWarnings("unchecked")
     @Test
     public void handleAddAuthenticationKey() throws Exception {
+        String password = "pass";
         IMappingServiceKey key = getDefualtKey();
-        MappingServiceValue value = getDefualtValue();
-        MappingEntry<MappingServiceValue> mappingEntry = new MappingEntry<MappingServiceValue>("value", value);
-        oneOf(lispDAO).get(weq(key));
-        oneOf(lispDAO).put(weq(key), weq(convertToArray(mappingEntry)));
-        assertEquals(true, testedMapServer.addAuthenticationKey(eid, key.getMask(), value.getKey()));
+        oneOf(lispDAO).put(weq(key),
+                weq((MappingEntry<String>[]) (Arrays.asList(new MappingEntry<String>(AbstractLispComponent.PASSWORD_SUBKEY, password)).toArray())));
+        testedMapServer.addAuthenticationKey(LispAFIConvertor.toContainer(eid), key.getMask(), password);
     }
 
     @Test
     public void handleGetAuthenticationKey() throws Exception {
         IMappingServiceKey key = getDefualtKey();
-        MappingServiceValue value = getDefualtValue();
-        Map<String, MappingServiceValue> map = new HashMap<String, MappingServiceValue>();
-        map.put("value", value);
-        atLeast(1).of(lispDAO).get(weq(key));
-        ret(map);
-        assertEquals(value.getKey(), testedMapServer.getAuthenticationKey(eid, key.getMask()));
+        oneOf(lispDAO).getSpecific(weq(key), with(AbstractLispComponent.PASSWORD_SUBKEY));
+        ret("password");
+        assertEquals("password", testedMapServer.getAuthenticationKey(LispAFIConvertor.toContainer(eid), key.getMask()));
     }
 
     @Test
@@ -304,76 +590,52 @@ public class MapServerTest extends BaseTestCase {
         testedMapServer.setShouldIterateMask(false);
         IMappingServiceKey key = getDefualtKey();
         IMappingServiceKey passKey = getKey(30);
-        MappingServiceValue value = getDefualtValue();
-        Map<String, MappingServiceValue> map = new HashMap<String, MappingServiceValue>();
-        map.put("value", value);
-        oneOf(lispDAO).get(weq(key));
-        allowing(lispDAO).get(weq(passKey));
-        ret(map);
-        assertEquals(null, testedMapServer.getAuthenticationKey(eid, key.getMask()));
+        oneOf(lispDAO).getSpecific(weq(key), with(AbstractLispComponent.PASSWORD_SUBKEY));
+        allowing(lispDAO).getSpecific(weq(passKey), with(AbstractLispComponent.PASSWORD_SUBKEY));
+        ret("password");
+        assertEquals(null, testedMapServer.getAuthenticationKey(LispAFIConvertor.toContainer(eid), key.getMask()));
     }
 
     @Test
     public void handleRemoveAuthenticationKey() throws Exception {
         IMappingServiceKey key = getDefualtKey();
-        MappingServiceValue value = new MappingServiceValue();
-        Map<String, MappingServiceValue> map = new HashMap<String, MappingServiceValue>();
-        map.put("value", value);
-        oneOf(lispDAO).get(weq(key));
-        ret(map);
-        oneOf(lispDAO).remove(weq(key));
-        assertEquals(true, testedMapServer.removeAuthenticationKey(eid, key.getMask()));
+        oneOf(lispDAO).removeSpecific(weq(key), with(AbstractLispComponent.PASSWORD_SUBKEY));
+        testedMapServer.removeAuthenticationKey(LispAFIConvertor.toContainer(eid), key.getMask());
     }
 
-    @Test
-    public void handleRemoveAuthenticationKeyWhereKeyDosntExist() throws Exception {
-        IMappingServiceKey key = getDefualtKey();
-        MappingServiceValue value = new MappingServiceValue();
-        Map<String, MappingServiceValue> map = new HashMap<String, MappingServiceValue>();
-        map.put("value", value);
-        oneOf(lispDAO).get(weq(key));
-        assertEquals(false, testedMapServer.removeAuthenticationKey(eid, key.getMask()));
+    private void addDefaultPutAndGetExpectations(LispAFIAddress addr, int mask) {
+        addPutExpectations(addr, mask);
+        addGetExpectations(addr, mask);
     }
 
-    @SuppressWarnings("rawtypes")
-    private MappingEntry[] convertToArray(MappingEntry<MappingServiceValue> entry) {
-        MappingEntry[] arr = new MappingEntry[1];
-        arr[0] = entry;
-        return arr;
+    private void addGetExpectations(LispAFIAddress address, int mask) {
+        addGetExpectations(address, mask, 0, 0, "password");
     }
 
-    private void addDAOExpectations(LispAddress address, int mask) {
-        addDAOEpectations(address, mask, 0, 0, "password");
+    private void addPutExpectations(LispAFIAddress address, int mask) {
+        oneOf(lispDAO).put(weq(MappingServiceKeyUtil.generateMappingServiceKey(LispAFIConvertor.toContainer(address), mask)),
+                with(mappingEntriesSaver));
     }
 
-    private void addDAOEpectations(LispAddress address, int mask, int withoutPassword, int withPassword, String password) {
-        LispAddress clonedAddress;
-        addPasswordExpextations(address, mask, withoutPassword, withPassword, password);
-        clonedAddress = (LispAddress) ((IMaskable) address).clone();
-        allowing(lispDAO).put(weq(MappingServiceKeyUtil.generateMappingServiceKey(clonedAddress, (byte) mask)), with(mappingEntriesSaver));
-    }
-
-    private void addPasswordExpextations(LispAddress address, int mask, int withoutPassword, int withPassword, String password) {
-        LispAddress clonedAddress = (LispAddress) ((IMaskable) address).clone();
+    private void addGetExpectations(LispAFIAddress address, int mask, int withoutPassword, int withPassword, String password) {
         if (withoutPassword > 0) {
-            Map<String, MappingServiceValue> result = new HashMap<String, MappingServiceValue>();
-            result.put("value", new MappingServiceValue());
-            allowing(lispDAO)
-                    .getSpecific(with(MappingServiceKeyUtil.generateMappingServiceKey(clonedAddress, (byte) withoutPassword)), with("value"));
+            String result = null;
+            result = null;
+            allowing(lispDAO).getSpecific(
+                    with(MappingServiceKeyUtil.generateMappingServiceKey(LispAFIConvertor.toContainer(address), withoutPassword)), with("password"));
             ret(result);
         }
-        clonedAddress = (LispAddress) ((IMaskable) address).clone();
         if (withPassword > 0) {
-            Map<String, MappingServiceValue> result = new HashMap<String, MappingServiceValue>();
-            result.put("value", new MappingServiceValue());
-            result.get("value").setKey(password);
-            allowing(lispDAO).get(with(MappingServiceKeyUtil.generateMappingServiceKey(clonedAddress, (byte) withPassword)));
+            String result = null;
+            result = password;
+            allowing(lispDAO).getSpecific(with(MappingServiceKeyUtil.generateMappingServiceKey(LispAFIConvertor.toContainer(address), withPassword)),
+                    with("password"));
             ret(result);
         }
         for (int i = mask; i >= 0; i--) {
-            clonedAddress = (LispAddress) ((IMaskable) address).clone();
-            allowing(lispDAO).get(with(MappingServiceKeyUtil.generateMappingServiceKey(clonedAddress, (byte) i)));
-            ret(new HashMap<String, MappingServiceValue>());
+            allowing(lispDAO).getSpecific(with(MappingServiceKeyUtil.generateMappingServiceKey(LispAFIConvertor.toContainer(address), i)),
+                    with("password"));
+            ret(null);
         }
     }
 
@@ -382,14 +644,41 @@ public class MapServerTest extends BaseTestCase {
     }
 
     private IMappingServiceKey getKey(int mask) {
-        IMappingServiceKey key = MappingServiceKeyUtil.generateMappingServiceKey(eid, mask);
+        IMappingServiceKey key = MappingServiceKeyUtil.generateMappingServiceKey(LispAFIConvertor.toContainer(eid), mask);
         return key;
     }
 
-    private MappingServiceValue getDefualtValue() {
-        MappingServiceValue value = new MappingServiceValue();
-        String password = "pass";
-        value.setKey(password);
-        return value;
+    private MapRegisterBuilder getDefaultMapRegisterBuilder() {
+        MapRegisterBuilder mrb = new MapRegisterBuilder();
+        mrb.setKeyId((short) 0);
+        mrb.setNonce((long) 0);
+        mrb.setWantMapNotify(false);
+        mrb.setProxyMapReply(false);
+        mrb.setEidToLocatorRecord(new ArrayList<EidToLocatorRecord>());
+        mrb.setAuthenticationData(new byte[0]);
+        return mrb;
+    }
+
+    private EidToLocatorRecordBuilder getDefaultEidToLocatorBuilder() {
+        EidToLocatorRecordBuilder builder = new EidToLocatorRecordBuilder();
+        builder.setAction(Action.NoAction);
+        builder.setAuthoritative(false);
+        builder.setLocatorRecord(new ArrayList<LocatorRecord>());
+        builder.setMapVersion((short) 0);
+        builder.setMaskLength((short) 32);
+        builder.setRecordTtl(60);
+        return builder;
+    }
+
+    private LocatorRecordBuilder getDefaultLocatorBuilder() {
+        LocatorRecordBuilder builder = new LocatorRecordBuilder();
+        builder.setLocalLocator(false);
+        builder.setMulticastPriority((short) 0);
+        builder.setMulticastWeight((short) 0);
+        builder.setPriority((short) 0);
+        builder.setRlocProbed(false);
+        builder.setRouted(false);
+        builder.setWeight((short) 0);
+        return builder;
     }
 }