From 07ec76b66365676efea1b96d541a20a7bdcb2e02 Mon Sep 17 00:00:00 2001
From: Blake Rouse <blake.rouse@elastic.co>
Date: Tue, 11 Apr 2023 11:41:05 -0400
Subject: [PATCH 1/2] Fix TestManager_FakeShipper flakiness.

---
 pkg/component/runtime/manager_test.go | 46 +++++++++++++++++++++++----
 1 file changed, 40 insertions(+), 6 deletions(-)

diff --git a/pkg/component/runtime/manager_test.go b/pkg/component/runtime/manager_test.go
index 8d0e808a860..262f97a62e6 100644
--- a/pkg/component/runtime/manager_test.go
+++ b/pkg/component/runtime/manager_test.go
@@ -2307,8 +2307,6 @@ func TestManager_FakeShipper(t *testing.T) {
 		5. Send `send_event` action to the component fake input (GRPC client); returns once sent.
 		6. Wait for `record_event` action to return from the shipper input (GRPC server).
 	*/
-	t.Skip("Flaky test: https://github.com/elastic/elastic-agent/issues/2301")
-
 	testPaths(t)
 
 	ctx, cancel := context.WithCancel(context.Background())
@@ -2412,14 +2410,21 @@ func TestManager_FakeShipper(t *testing.T) {
 	defer subCancel()
 	subErrCh := make(chan error)
 	go func() {
-		shipperOn := false
+		shipperInputOn := false
+		shipperOutputOn := false
 		compConnected := false
+		eventSent := false
 
 		sendEvent := func() (bool, error) {
-			if !shipperOn || !compConnected {
+			if !shipperInputOn || !shipperOutputOn || !compConnected {
 				// wait until connected
 				return false, nil
 			}
+			if eventSent {
+				// other path already sent event
+				return false, nil
+			}
+			eventSent = true
 
 			// send an event between component and the fake shipper
 			eventID, err := uuid.NewV4()
@@ -2471,7 +2476,7 @@ func TestManager_FakeShipper(t *testing.T) {
 						if unit.State == client.UnitStateFailed {
 							subErrCh <- fmt.Errorf("unit failed: %s", unit.Message)
 						} else if unit.State == client.UnitStateHealthy {
-							shipperOn = true
+							shipperInputOn = true
 							ok, err := sendEvent()
 							if ok {
 								if err != nil {
@@ -2493,7 +2498,36 @@ func TestManager_FakeShipper(t *testing.T) {
 							subErrCh <- fmt.Errorf("unit reported unexpected state: %v", unit.State)
 						}
 					} else {
-						subErrCh <- errors.New("unit missing: fake-input")
+						subErrCh <- errors.New("input unit missing: fake-input")
+					}
+					unit, ok = state.Units[ComponentUnitKey{UnitType: client.UnitTypeOutput, UnitID: "fake-default"}]
+					if ok {
+						if unit.State == client.UnitStateFailed {
+							subErrCh <- fmt.Errorf("unit failed: %s", unit.Message)
+						} else if unit.State == client.UnitStateHealthy {
+							shipperOutputOn = true
+							ok, err := sendEvent()
+							if ok {
+								if err != nil {
+									subErrCh <- err
+								} else {
+									// successful; turn it all off
+									err := m.Update([]component.Component{})
+									if err != nil {
+										subErrCh <- err
+									}
+								}
+							}
+						} else if unit.State == client.UnitStateStopped {
+							subErrCh <- nil
+						} else if unit.State == client.UnitStateStarting {
+							// acceptable
+						} else {
+							// unknown state that should not have occurred
+							subErrCh <- fmt.Errorf("unit reported unexpected state: %v", unit.State)
+						}
+					} else {
+						subErrCh <- errors.New("output unit missing: fake-input")
 					}
 				}
 			case state := <-compSub.Ch():

From aab68e068f1f8887cdf747e3a750c7cfa462b7ab Mon Sep 17 00:00:00 2001
From: Blake Rouse <blake.rouse@elastic.co>
Date: Wed, 12 Apr 2023 15:47:16 -0400
Subject: [PATCH 2/2] Fix bad error message in test.

---
 pkg/component/runtime/manager_test.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pkg/component/runtime/manager_test.go b/pkg/component/runtime/manager_test.go
index 262f97a62e6..3ab469ea264 100644
--- a/pkg/component/runtime/manager_test.go
+++ b/pkg/component/runtime/manager_test.go
@@ -2498,7 +2498,7 @@ func TestManager_FakeShipper(t *testing.T) {
 							subErrCh <- fmt.Errorf("unit reported unexpected state: %v", unit.State)
 						}
 					} else {
-						subErrCh <- errors.New("input unit missing: fake-input")
+						subErrCh <- errors.New("input unit missing: fake-default")
 					}
 					unit, ok = state.Units[ComponentUnitKey{UnitType: client.UnitTypeOutput, UnitID: "fake-default"}]
 					if ok {
@@ -2527,7 +2527,7 @@ func TestManager_FakeShipper(t *testing.T) {
 							subErrCh <- fmt.Errorf("unit reported unexpected state: %v", unit.State)
 						}
 					} else {
-						subErrCh <- errors.New("output unit missing: fake-input")
+						subErrCh <- errors.New("output unit missing: fake-default")
 					}
 				}
 			case state := <-compSub.Ch():