diff --git a/internal/bbsim/devices/onu.go b/internal/bbsim/devices/onu.go
index 5eed730..1353a50 100644
--- a/internal/bbsim/devices/onu.go
+++ b/internal/bbsim/devices/onu.go
@@ -382,7 +382,7 @@
 	// clean the ONU state
 	o.Flows = []FlowKey{}
 	o.PonPort.removeOnuId(o.ID)
-	o.PonPort.removeAllocIdsForOnuSn(o.SerialNumber)
+	o.PonPort.removeAllocId(o.SerialNumber)
 	o.PonPort.removeGemPortBySn(o.SerialNumber)
 
 	o.onuAlarmsInfoLock.Lock()
@@ -807,7 +807,7 @@
 			o.MibDataSync = 0
 
 			// if the MIB reset is successful then remove all the stored AllocIds and GemPorts
-			o.PonPort.removeAllocIdsForOnuSn(o.SerialNumber)
+			o.PonPort.removeAllocId(o.SerialNumber)
 			o.PonPort.removeGemPortBySn(o.SerialNumber)
 		}
 	case omci.MibUploadRequestType:
@@ -887,7 +887,6 @@
 			}).Debug("set-onu-admin-lock-state")
 		case me.TContClassID:
 			allocId := msgObj.Attributes[me.TCont_AllocId].(uint16)
-			entityID := msgObj.Attributes[me.ManagedEntityID].(uint16)
 
 			// if the AllocId is 255 (0xFF) or 65535 (0xFFFF) it means we are removing it,
 			// otherwise we are adding it
@@ -899,15 +898,15 @@
 					"AllocId":      allocId,
 					"SerialNumber": o.Sn(),
 				}).Trace("freeing-alloc-id-via-omci")
-				o.PonPort.removeAllocId(entityID)
+				o.PonPort.removeAllocId(o.SerialNumber)
 			} else {
-				if used, allocObj := o.PonPort.isAllocIdAllocated(entityID); used {
+				if used, sn := o.PonPort.isAllocIdAllocated(allocId); used {
 					onuLogger.WithFields(log.Fields{
 						"IntfId":       o.PonPortID,
 						"OnuId":        o.ID,
 						"AllocId":      allocId,
 						"SerialNumber": o.Sn(),
-					}).Errorf("allocid-already-allocated-to-onu-with-sn-%s", common.OnuSnToString(allocObj.OnuSn))
+					}).Errorf("allocid-already-allocated-to-onu-with-sn-%s", common.OnuSnToString(sn))
 					success = false
 				} else {
 					onuLogger.WithFields(log.Fields{
@@ -917,7 +916,7 @@
 						"AllocId":      allocId,
 						"SerialNumber": o.Sn(),
 					}).Trace("storing-alloc-id-via-omci")
-					o.PonPort.storeAllocId(entityID, allocId, o.SerialNumber)
+					o.PonPort.storeAllocId(allocId, o.SerialNumber)
 				}
 			}
 		case me.EthernetFrameExtendedPmClassID,