diff --git a/internal/bbsim/devices/onu.go b/internal/bbsim/devices/onu.go
index af135de..3ee5a7f 100644
--- a/internal/bbsim/devices/onu.go
+++ b/internal/bbsim/devices/onu.go
@@ -197,9 +197,9 @@
 			// NOTE should disabled state be different for oper_disabled (emulating an error) and admin_disabled (received a disabled call via VOLTHA)?
 			{Name: OnuTxDisable, Src: []string{OnuStateEnabled, OnuStatePonDisabled, OnuStateImageActivated, OnuStateImageDownloadError, OnuStateImageCommitted}, Dst: OnuStateDisabled},
 			// ONU state when PON port is disabled but ONU is power ON(more states should be added in src?)
-			{Name: OnuTxPonDisable, Src: []string{OnuStateEnabled, OnuStateImageActivated, OnuStateImageDownloadError, OnuStateImageCommitted}, Dst: OnuStatePonDisabled},
+			{Name: OnuTxPonDisable, Src: []string{OnuStateEnabled, OnuStateImageActivated, OnuStateImageDownloadError, OnuStateImageCommitted, OnuStateImageDownloadComplete}, Dst: OnuStatePonDisabled},
 			// Software Image Download related states
-			{Name: OnuTxStartImageDownload, Src: []string{OnuStateEnabled, OnuStateImageDownloadComplete, OnuStateImageDownloadError}, Dst: OnuStateImageDownloadStarted},
+			{Name: OnuTxStartImageDownload, Src: []string{OnuStateEnabled, OnuStateImageDownloadComplete, OnuStateImageDownloadError, OnuStateImageCommitted}, Dst: OnuStateImageDownloadStarted},
 			{Name: OnuTxProgressImageDownload, Src: []string{OnuStateImageDownloadStarted}, Dst: OnuStateImageDownloadInProgress},
 			{Name: OnuTxCompleteImageDownload, Src: []string{OnuStateImageDownloadInProgress}, Dst: OnuStateImageDownloadComplete},
 			{Name: OnuTxFailImageDownload, Src: []string{OnuStateImageDownloadInProgress}, Dst: OnuStateImageDownloadError},
diff --git a/internal/bbsim/devices/onu_state_machine_test.go b/internal/bbsim/devices/onu_state_machine_test.go
index 1de9eeb..5784dd7 100644
--- a/internal/bbsim/devices/onu_state_machine_test.go
+++ b/internal/bbsim/devices/onu_state_machine_test.go
@@ -26,10 +26,10 @@
 
 func Test_Onu_StateMachine_enable(t *testing.T) {
 	onu := createTestOnu()
-	assert.Equal(t, onu.InternalState.Current(), "initialized")
-	_ = onu.InternalState.Event("discover")
-	assert.Equal(t, onu.InternalState.Current(), "discovered")
-	_ = onu.InternalState.Event("enable")
+	assert.Equal(t, onu.InternalState.Current(), OnuStateInitialized)
+	_ = onu.InternalState.Event(OnuTxDiscover)
+	assert.Equal(t, onu.InternalState.Current(), OnuStateDiscovered)
+	_ = onu.InternalState.Event(OnuTxEnable)
 	assert.Equal(t, onu.InternalState.Current(), OnuStateEnabled)
 }
 
@@ -63,6 +63,55 @@
 	assert.Equal(t, len(onu.PonPort.AllocatedGemPorts), 0)
 }
 
+func Test_Onu_StateMachine_pon_disable(t *testing.T) {
+	onu := createTestOnu()
+	var err error
+
+	onu.InternalState.SetState(OnuStateEnabled)
+	err = onu.InternalState.Event(OnuTxPonDisable)
+	assert.NilError(t, err)
+	assert.Equal(t, OnuStatePonDisabled, onu.InternalState.Current())
+
+	onu.InternalState.SetState(OnuStateImageDownloadComplete)
+	err = onu.InternalState.Event(OnuTxPonDisable)
+	assert.NilError(t, err)
+	assert.Equal(t, OnuStatePonDisabled, onu.InternalState.Current())
+}
+
+func Test_Onu_StateMachine_software_image(t *testing.T) {
+	onu := createTestOnu()
+	var err error
+
+	// happy path
+	onu.InternalState.SetState(OnuStateEnabled)
+	err = onu.InternalState.Event(OnuTxStartImageDownload)
+	assert.NilError(t, err)
+	assert.Equal(t, OnuStateImageDownloadStarted, onu.InternalState.Current())
+
+	err = onu.InternalState.Event(OnuTxProgressImageDownload)
+	assert.NilError(t, err)
+	assert.Equal(t, OnuStateImageDownloadInProgress, onu.InternalState.Current())
+
+	err = onu.InternalState.Event(OnuTxCompleteImageDownload)
+	assert.NilError(t, err)
+	assert.Equal(t, OnuStateImageDownloadComplete, onu.InternalState.Current())
+
+	err = onu.InternalState.Event(OnuTxActivateImage)
+	assert.NilError(t, err)
+	assert.Equal(t, OnuStateImageActivated, onu.InternalState.Current())
+
+	// after image activate we get an ONU reboot, thus the state is back to Enabled before committing
+	onu.InternalState.SetState(OnuStateEnabled)
+	err = onu.InternalState.Event(OnuTxCommitImage)
+	assert.NilError(t, err)
+	assert.Equal(t, OnuStateImageCommitted, onu.InternalState.Current())
+
+	// but we should be able to start a new download
+	err = onu.InternalState.Event(OnuTxStartImageDownload)
+	assert.NilError(t, err)
+	assert.Equal(t, OnuStateImageDownloadStarted, onu.InternalState.Current())
+}
+
 // check that I can go to auth_started only if
 // - the GemPort is set
 // - the eapolFlow is received