From 5bbcff1f9a267b86553ad60d09a3cab94395c9f0 Mon Sep 17 00:00:00 2001
From: Wesley GALIPO <wesley.galipo@ovhcloud.com>
Date: Mon, 22 May 2023 09:36:00 +0000
Subject: [PATCH] add failure_or_done notification stragegy

Signed-off-by: Wesley GALIPO <wesley.galipo@ovhcloud.com>
---
 pkg/notify/init/init.go |  4 ++++
 pkg/notify/messages.go  | 11 +++++++++++
 utask.go                |  2 ++
 3 files changed, 17 insertions(+)

diff --git a/pkg/notify/init/init.go b/pkg/notify/init/init.go
index bff71ba2..52b0ac67 100644
--- a/pkg/notify/init/init.go
+++ b/pkg/notify/init/init.go
@@ -148,6 +148,10 @@ func validateStrategyForAction(action, strategy string) error {
 		if action == notify.TaskValidationKey {
 			return errNotAllowed
 		}
+	case utask.NotificationStrategyFailureOrDone:
+		if action == notify.TaskValidationKey {
+			return errNotAllowed
+		}
 	default:
 		return errUnknown
 	}
diff --git a/pkg/notify/messages.go b/pkg/notify/messages.go
index f8f3d8a9..8950db55 100644
--- a/pkg/notify/messages.go
+++ b/pkg/notify/messages.go
@@ -236,6 +236,10 @@ func checkIfDeliverMessageFromTaskState(m *Message, strategy string) bool {
 		if v, ok := m.Fields["state"]; ok && v == stateBlocked {
 			send = true
 		}
+	case utask.NotificationStrategyFailureOrDone:
+		if v, ok := m.Fields["state"]; ok && (v == stateBlocked || v == step.StateDone) {
+			send = true
+		}
 	case utask.NotificationStrategySilent:
 	}
 
@@ -254,6 +258,13 @@ func checkIfDeliverMessageFromStepState(m *Message, strategy string) bool {
 				send = true
 			}
 		}
+	case utask.NotificationStrategyFailureOrDone:
+		if v, ok := m.Fields["step_state"]; ok {
+			switch v {
+			case step.StateFatalError, step.StateCrashed, step.StateAfterrunError, step.StateDone:
+				send = true
+			}
+		}
 	case utask.NotificationStrategySilent:
 	}
 
diff --git a/utask.go b/utask.go
index f296312a..d1dcfe20 100644
--- a/utask.go
+++ b/utask.go
@@ -100,6 +100,8 @@ const (
 	NotificationStrategyAlways = "always"
 	// NotificationStrategyFailureOnly corresponds to the mode where notifications will only be sent if the state is BLOCKED
 	NotificationStrategyFailureOnly = "failure_only"
+	// NotificationStrategyFailureOrDone corresponds to the mode where notifications will only be sent if the state is BLOCKED or DONE
+	NotificationStrategyFailureOrDone = "failure_or_done"
 
 	// UsernamesSeparator corresponds to the separator used to break a string into a list of usernames and vice versa.
 	UsernamesSeparator = ","