From 45bf1fd887c5e7dd84227c27998e55d1c1c70397 Mon Sep 17 00:00:00 2001
From: Lukasz <120112546+lukaszcl@users.noreply.github.com>
Date: Fri, 12 Jan 2024 14:53:08 +0100
Subject: [PATCH] Allow custom hooks for ClNode in e2e tests (#11745)

* Allow custom hooks for ClNode in e2e tests

* Fix for json.Marshall
---
 integration-tests/docker/test_env/cl_node.go | 44 +++++++++++++-------
 1 file changed, 28 insertions(+), 16 deletions(-)

diff --git a/integration-tests/docker/test_env/cl_node.go b/integration-tests/docker/test_env/cl_node.go
index a00dea3a35a..1466d03c6f1 100644
--- a/integration-tests/docker/test_env/cl_node.go
+++ b/integration-tests/docker/test_env/cl_node.go
@@ -51,7 +51,10 @@ type ClNode struct {
 	PostgresDb            *test_env.PostgresDb    `json:"postgresDb"`
 	UserEmail             string                  `json:"userEmail"`
 	UserPassword          string                  `json:"userPassword"`
-	AlwaysPullImage       bool
+	AlwaysPullImage       bool                    `json:"-"`
+	PostStartsHooks       []tc.ContainerHook      `json:"-"`
+	PostStopsHooks        []tc.ContainerHook      `json:"-"`
+	PreTerminatesHooks    []tc.ContainerHook      `json:"-"`
 	t                     *testing.T
 	l                     zerolog.Logger
 	ls                    *logstream.LogStream
@@ -140,12 +143,32 @@ func NewClNode(networks []string, imageName, imageVersion string, nodeConfig *ch
 		PostgresDb:   pgDb,
 		l:            log.Logger,
 	}
+	n.SetDefaultHooks()
 	for _, opt := range opts {
 		opt(n)
 	}
 	return n, nil
 }
 
+func (n *ClNode) SetDefaultHooks() {
+	n.PostStartsHooks = []tc.ContainerHook{
+		func(ctx context.Context, c tc.Container) error {
+			if n.ls != nil {
+				return n.ls.ConnectContainer(ctx, c, "cl-node")
+			}
+			return nil
+		},
+	}
+	n.PostStopsHooks = []tc.ContainerHook{
+		func(ctx context.Context, c tc.Container) error {
+			if n.ls != nil {
+				return n.ls.DisconnectContainer(c)
+			}
+			return nil
+		},
+	}
+}
+
 func (n *ClNode) SetTestLogger(t *testing.T) {
 	n.l = logging.GetTestLogger(t)
 	n.t = t
@@ -475,22 +498,11 @@ func (n *ClNode) getContainerRequest(secrets string) (
 			},
 		},
 		LifecycleHooks: []tc.ContainerLifecycleHooks{
-			{PostStarts: []tc.ContainerHook{
-				func(ctx context.Context, c tc.Container) error {
-					if n.ls != nil {
-						return n.ls.ConnectContainer(ctx, c, "cl-node")
-					}
-					return nil
-				},
+			{
+				PostStarts:    n.PostStartsHooks,
+				PostStops:     n.PostStopsHooks,
+				PreTerminates: n.PreTerminatesHooks,
 			},
-				PostStops: []tc.ContainerHook{
-					func(ctx context.Context, c tc.Container) error {
-						if n.ls != nil {
-							return n.ls.DisconnectContainer(c)
-						}
-						return nil
-					},
-				}},
 		},
 	}, nil
 }