Merge "[VOL-3286] Sending PortStats only when OLT Enabled"
diff --git a/internal/bbsim/devices/olt.go b/internal/bbsim/devices/olt.go
index 2799de8..e0d1556 100644
--- a/internal/bbsim/devices/olt.go
+++ b/internal/bbsim/devices/olt.go
@@ -902,25 +902,9 @@
 		}).Infof("Failed to transition ONU to disabled state: %s", err.Error())
 	}
 
-	time.Sleep(1 * time.Second)
-
 	// ONU Re-Discovery
 	if o.InternalState.Current() == "enabled" && pon.InternalState.Current() == "enabled" {
-		if err := _onu.InternalState.Event("initialize"); err != nil {
-			oltLogger.WithFields(log.Fields{
-				"IntfId": _onu.PonPortID,
-				"OnuSn":  _onu.Sn(),
-				"OnuId":  _onu.ID,
-			}).Infof("Failed to transition ONU to initialized state: %s", err.Error())
-		}
-
-		if err := _onu.InternalState.Event("discover"); err != nil {
-			oltLogger.WithFields(log.Fields{
-				"IntfId": _onu.PonPortID,
-				"OnuSn":  _onu.Sn(),
-				"OnuId":  _onu.ID,
-			}).Infof("Failed to transition ONU to discovered state: %s", err.Error())
-		}
+		go _onu.ReDiscoverOnu()
 	}
 
 	return new(openolt.Empty), nil
diff --git a/internal/bbsim/devices/onu.go b/internal/bbsim/devices/onu.go
index a21d96a..beb3779 100644
--- a/internal/bbsim/devices/onu.go
+++ b/internal/bbsim/devices/onu.go
@@ -245,8 +245,12 @@
 					},
 				}
 				o.Channel <- msg
+
+				// verify all the flows removes are handled and
 				// terminate the ONU's ProcessOnuMessages Go routine
-				close(o.Channel)
+				if len(o.FlowIds) == 0 {
+					close(o.Channel)
+				}
 			},
 			"before_start_auth": func(e *fsm.Event) {
 				if o.EapolFlowReceived == false {
@@ -896,6 +900,12 @@
 		// so that we can properly set these two flag when the flow is removed
 		o.EapolFlowReceived = false
 		o.DhcpFlowReceived = false
+
+		// check if ONU delete is performed and
+		// terminate the ONU's ProcessOnuMessages Go routine
+		if o.InternalState.Current() == "disabled" {
+			close(o.Channel)
+		}
 	}
 }
 
@@ -1120,3 +1130,31 @@
 		}
 	}
 }
+
+func (onu *Onu) ReDiscoverOnu() {
+	// Wait for few seconds to be sure of the cleanup
+	time.Sleep(5 * time.Second)
+
+	onuLogger.WithFields(log.Fields{
+		"IntfId": onu.PonPortID,
+		"OnuId":  onu.ID,
+		"OnuSn":  onu.Sn(),
+	}).Debug("Send ONU Re-Discovery")
+
+	// ONU Re-Discovery
+	if err := onu.InternalState.Event("initialize"); err != nil {
+		log.WithFields(log.Fields{
+			"IntfId": onu.PonPortID,
+			"OnuSn":  onu.Sn(),
+			"OnuId":  onu.ID,
+		}).Infof("Failed to transition ONU to initialized state: %s", err.Error())
+	}
+
+	if err := onu.InternalState.Event("discover"); err != nil {
+		log.WithFields(log.Fields{
+			"IntfId": onu.PonPortID,
+			"OnuSn":  onu.Sn(),
+			"OnuId":  onu.ID,
+		}).Infof("Failed to transition ONU to discovered state: %s", err.Error())
+	}
+}
diff --git a/internal/bbsim/devices/onu_flow_test.go b/internal/bbsim/devices/onu_flow_test.go
index b01859e..e1c801f 100644
--- a/internal/bbsim/devices/onu_flow_test.go
+++ b/internal/bbsim/devices/onu_flow_test.go
@@ -510,6 +510,15 @@
 // checks that when the last flow is removed we reset the stored flags in the ONU
 func Test_HandleFlowRemoveFlowId_LastFlow(t *testing.T) {
 	onu := createMockOnu(1, 1, 900, 900, true, false)
+
+	onu.InternalState = fsm.NewFSM(
+		"enabled",
+		fsm.Events{
+			{Name: "disable", Src: []string{"enabled"}, Dst: "disabled"},
+		},
+		fsm.Callbacks{},
+	)
+
 	onu.GemPortAdded = true
 	onu.DhcpFlowReceived = true
 	onu.EapolFlowReceived = true
@@ -530,4 +539,4 @@
 	assert.Equal(t, onu.GemPortAdded, false)
 	assert.Equal(t, onu.DhcpFlowReceived, false)
 	assert.Equal(t, onu.EapolFlowReceived, false)
-}
\ No newline at end of file
+}