Skip to content

Commit

Permalink
fix(notifications): title customization
Browse files Browse the repository at this point in the history
  • Loading branch information
@piksel
piksel committed Feb 11, 2022
1 parent 70ca4c2 commit 494859c
Show file tree
Hide file tree
Showing 5 changed files with 132 additions and 46 deletions.
10 changes: 10 additions & 0 deletions internal/flags/flags.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -320,6 +320,16 @@ Should only be used for testing.`)
viper.GetBool("WATCHTOWER_NOTIFICATION_REPORT"),
"Use the session report as the notification template data")

flags.StringP(
"notification-title-tag",
"",
viper.GetString("WATCHTOWER_NOTIFICATION_TITLE_TAG"),
"Title prefix tag for notifications")

flags.Bool("notification-skip-title",
viper.GetBool("WATCHTOWER_NOTIFICATION_SKIP_TITLE"),
"Do not pass the title param to notifications")

flags.String(
"warn-on-head-failure",
viper.GetString("WATCHTOWER_WARN_ON_HEAD_FAILURE"),
Expand Down
60 changes: 44 additions & 16 deletions pkg/notifications/notifier.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package notifications

import (
"os"
"strings"
"time"

ty "github.com/containrrr/watchtower/pkg/types"
Expand All @@ -10,6 +11,11 @@ import (
"github.com/spf13/cobra"
)

type TitleProvider struct {
Tag string
Hostname string
}

// NewNotifier creates and returns a new Notifier, using global configuration.
func NewNotifier(c *cobra.Command) ty.Notifier {
f := c.PersistentFlags()
Expand All @@ -30,10 +36,10 @@ func NewNotifier(c *cobra.Command) ty.Notifier {
tplString, _ := f.GetString("notification-template")
urls, _ := f.GetStringArray("notification-url")

hostname := GetHostname(c)
urls, delay := AppendLegacyUrls(urls, c, GetTitle(hostname))
data := GetTemplateData(c)
urls, delay := AppendLegacyUrls(urls, c, data.Title)

return newShoutrrrNotifier(tplString, acceptedLogLevels, !reportTemplate, hostname, delay, urls...)
return newShoutrrrNotifier(tplString, acceptedLogLevels, !reportTemplate, data, delay, urls...)
}

// AppendLegacyUrls creates shoutrrr equivalent URLs from legacy notification flags
Expand Down Expand Up @@ -84,28 +90,50 @@ func AppendLegacyUrls(urls []string, cmd *cobra.Command, title string) ([]string
return urls, delay
}

// GetTitle returns a common notification title with hostname appended
func GetTitle(hostname string) string {
title := "Watchtower updates"
// GetTitle formats the title based on the passed hostname and tag
func GetTitle(hostname string, tag string) string {
tb := strings.Builder{}

if tag != "" {
tb.WriteRune('[')
tb.WriteString(tag)
tb.WriteRune(']')
tb.WriteRune(' ')
}

tb.WriteString("Watchtower updates")

if hostname != "" {
title += " on " + hostname
tb.WriteString(" on ")
tb.WriteString(hostname)
}
return title
}

// GetHostname returns the hostname as set by args or resolved from OS
func GetHostname(c *cobra.Command) string {
return tb.String()
}

// GetTemplateData populates the static notification data from flags and environment
func GetTemplateData(c *cobra.Command) StaticData {
f := c.PersistentFlags()

hostname, _ := f.GetString("notifications-hostname")
if hostname == "" {
hostname, _ = os.Hostname()
}

if hostname != "" {
return hostname
} else if hostname, err := os.Hostname(); err == nil {
return hostname
title := ""
if skip, _ := f.GetBool("notification-skip-title"); !skip {
tag, _ := f.GetString("notification-title-tag")
if tag == "" {
// For legacy email support
tag, _ = f.GetString("notification-email-subjecttag")
}
title = GetTitle(hostname, tag)
}

return ""
return StaticData{
Host: hostname,
Title: title,
}
}

// ColorHex is the default notification color used for services that support it (formatted as a CSS hex string)
Expand Down
68 changes: 54 additions & 14 deletions pkg/notifications/notifier_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,17 +38,58 @@ var _ = Describe("notifications", func() {
"test.host",
})
Expect(err).NotTo(HaveOccurred())
hostname := notifications.GetHostname(command)
title := notifications.GetTitle(hostname)
data := notifications.GetTemplateData(command)
title := data.Title
Expect(title).To(Equal("Watchtower updates on test.host"))
})
})
When("no hostname can be resolved", func() {
It("should use the default simple title", func() {
title := notifications.GetTitle("")
title := notifications.GetTitle("", "")
Expect(title).To(Equal("Watchtower updates"))
})
})
When("title tag is set", func() {
It("should use the prefix in the title", func() {
command := cmd.NewRootCommand()
flags.RegisterNotificationFlags(command)

Expect(command.ParseFlags([]string{
"--notification-title-tag",
"PREFIX",
})).To(Succeed())

data := notifications.GetTemplateData(command)
Expect(data.Title).To(HavePrefix("[PREFIX]"))
})
})
When("legacy email tag is set", func() {
It("should use the prefix in the title", func() {
command := cmd.NewRootCommand()
flags.RegisterNotificationFlags(command)

Expect(command.ParseFlags([]string{
"--notification-email-subjecttag",
"PREFIX",
})).To(Succeed())

data := notifications.GetTemplateData(command)
Expect(data.Title).To(HavePrefix("[PREFIX]"))
})
})
When("the skip title flag is set", func() {
It("should return an empty title", func() {
command := cmd.NewRootCommand()
flags.RegisterNotificationFlags(command)

Expect(command.ParseFlags([]string{
"--notification-skip-title",
})).To(Succeed())

data := notifications.GetTemplateData(command)
Expect(data.Title).To(BeEmpty())
})
})
})
Describe("the slack notifier", func() {
// builderFn := notifications.NewSlackNotifier
Expand All @@ -60,8 +101,8 @@ var _ = Describe("notifications", func() {
channel := "123456789"
token := "abvsihdbau"
color := notifications.ColorInt
hostname := notifications.GetHostname(command)
title := url.QueryEscape(notifications.GetTitle(hostname))
data := notifications.GetTemplateData(command)
title := url.QueryEscape(data.Title)
expected := fmt.Sprintf("discord://%s@%s?color=0x%x&colordebug=0x0&colorerror=0x0&colorinfo=0x0&colorwarn=0x0&title=%s&username=watchtower", token, channel, color, title)
buildArgs := func(url string) []string {
return []string{
Expand Down Expand Up @@ -89,8 +130,8 @@ var _ = Describe("notifications", func() {
tokenB := "BBBBBBBBB"
tokenC := "123456789123456789123456"
color := url.QueryEscape(notifications.ColorHex)
hostname := notifications.GetHostname(command)
title := url.QueryEscape(notifications.GetTitle(hostname))
data := notifications.GetTemplateData(command)
title := url.QueryEscape(data.Title)
iconURL := "https://containrrr.dev/watchtower-sq180.png"
iconEmoji := "whale"

Expand Down Expand Up @@ -145,8 +186,8 @@ var _ = Describe("notifications", func() {

token := "aaa"
host := "shoutrrr.local"
hostname := notifications.GetHostname(command)
title := url.QueryEscape(notifications.GetTitle(hostname))
data := notifications.GetTemplateData(command)
title := url.QueryEscape(data.Title)

expectedOutput := fmt.Sprintf("gotify://%s/%s?title=%s", host, token, title)

Expand Down Expand Up @@ -174,8 +215,8 @@ var _ = Describe("notifications", func() {
tokenB := "33333333012222222222333333333344"
tokenC := "44444444-4444-4444-8444-cccccccccccc"
color := url.QueryEscape(notifications.ColorHex)
hostname := notifications.GetHostname(command)
title := url.QueryEscape(notifications.GetTitle(hostname))
data := notifications.GetTemplateData(command)
title := url.QueryEscape(data.Title)

hookURL := fmt.Sprintf("https://outlook.office.com/webhook/%s/IncomingWebhook/%s/%s", tokenA, tokenB, tokenC)
expectedOutput := fmt.Sprintf("teams://%s/%s/%s?color=%s&title=%s", tokenA, tokenB, tokenC, color, title)
Expand Down Expand Up @@ -266,9 +307,8 @@ func testURL(args []string, expectedURL string) {
err := command.ParseFlags(args)
Expect(err).NotTo(HaveOccurred())

hostname := notifications.GetHostname(command)
title := notifications.GetTitle(hostname)
urls, _ := notifications.AppendLegacyUrls([]string{}, command, title)
data := notifications.GetTemplateData(command)
urls, _ := notifications.AppendLegacyUrls([]string{}, command, data.Title)

Expect(err).NotTo(HaveOccurred())

Expand Down
27 changes: 16 additions & 11 deletions pkg/notifications/shoutrrr.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,7 @@ type shoutrrrTypeNotifier struct {
done chan bool
legacyTemplate bool
params *types.Params
hostname string
data StaticData
}

// GetScheme returns the scheme part of a Shoutrrr URL
Expand All @@ -79,11 +79,9 @@ func (n *shoutrrrTypeNotifier) GetNames() []string {
return names
}

func newShoutrrrNotifier(tplString string, acceptedLogLevels []log.Level, legacy bool, hostname string, delay time.Duration, urls ...string) t.Notifier {
func newShoutrrrNotifier(tplString string, acceptedLogLevels []log.Level, legacy bool, data StaticData, delay time.Duration, urls ...string) t.Notifier {

notifier := createNotifier(urls, acceptedLogLevels, tplString, legacy)
notifier.hostname = hostname
notifier.params = &types.Params{"title": GetTitle(hostname)}
notifier := createNotifier(urls, acceptedLogLevels, tplString, legacy, data)
log.AddHook(notifier)

// Do the sending in a separate goroutine so we don't block the main process.
Expand All @@ -92,7 +90,7 @@ func newShoutrrrNotifier(tplString string, acceptedLogLevels []log.Level, legacy
return notifier
}

func createNotifier(urls []string, levels []log.Level, tplString string, legacy bool) *shoutrrrTypeNotifier {
func createNotifier(urls []string, levels []log.Level, tplString string, legacy bool, data StaticData) *shoutrrrTypeNotifier {
tpl, err := getShoutrrrTemplate(tplString, legacy)
if err != nil {
log.Errorf("Could not use configured notification template: %s. Using default template", err)
Expand All @@ -112,6 +110,10 @@ func createNotifier(urls []string, levels []log.Level, tplString string, legacy
logLevels: levels,
template: tpl,
legacyTemplate: legacy,
data: data,
params: &types.Params{
"title": data.Title,
},
}
}

Expand Down Expand Up @@ -149,9 +151,7 @@ func (n *shoutrrrTypeNotifier) buildMessage(data Data) (string, error) {
}

func (n *shoutrrrTypeNotifier) sendEntries(entries []*log.Entry, report t.Report) {
title, _ := n.params.Title()
host := n.hostname
msg, err := n.buildMessage(Data{entries, report, title, host})
msg, err := n.buildMessage(Data{n.data, entries, report})

if msg == "" {
// Log in go func in case we entered from Fire to avoid stalling
Expand Down Expand Up @@ -240,10 +240,15 @@ func getShoutrrrTemplate(tplString string, legacy bool) (tpl *template.Template,
return
}

// StaticData is the part of the notification template data model set upon initialization
type StaticData struct {
Title string
Host string
}

// Data is the notification template data model
type Data struct {
StaticData
Entries []*log.Entry
Report t.Report
Title string
Host string
}
13 changes: 8 additions & 5 deletions pkg/notifications/shoutrrr_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,11 +49,14 @@ var mockDataAllFresh = Data{

func mockDataFromStates(states ...s.State) Data {
hostname := "Mock"
prefix := ""
return Data{
Entries: legacyMockData.Entries,
Report: mocks.CreateMockProgressReport(states...),
Title: GetTitle(hostname),
Host: hostname,
StaticData: StaticData{
Title: GetTitle(hostname, prefix),
Host: hostname,
},
}
}

Expand All @@ -77,7 +80,7 @@ var _ = Describe("Shoutrrr", func() {
cmd := new(cobra.Command)
flags.RegisterNotificationFlags(cmd)

shoutrrr := createNotifier([]string{}, logrus.AllLevels, "", true)
shoutrrr := createNotifier([]string{}, logrus.AllLevels, "", true, StaticData{})

entries := []*logrus.Entry{
{
Expand Down Expand Up @@ -233,15 +236,15 @@ Turns out everything is on fire
When("batching notifications", func() {
When("no messages are queued", func() {
It("should not send any notification", func() {
shoutrrr := newShoutrrrNotifier("", allButTrace, true, "", time.Duration(0), "logger://")
shoutrrr := newShoutrrrNotifier("", allButTrace, true, StaticData{}, time.Duration(0), "logger://")
shoutrrr.StartNotification()
shoutrrr.SendNotification(nil)
Consistently(logBuffer).ShouldNot(gbytes.Say(`Shoutrrr:`))
})
})
When("at least one message is queued", func() {
It("should send a notification", func() {
shoutrrr := newShoutrrrNotifier("", allButTrace, true, "", time.Duration(0), "logger://")
shoutrrr := newShoutrrrNotifier("", allButTrace, true, StaticData{}, time.Duration(0), "logger://")
shoutrrr.StartNotification()
logrus.Info("This log message is sponsored by ContainrrrVPN")
shoutrrr.SendNotification(nil)
Expand Down

0 comments on commit 494859c

Please sign in to comment.