-
Notifications
You must be signed in to change notification settings - Fork 1.3k
Expand file tree
/
Copy pathexp_scaletest_smtp.go
More file actions
112 lines (99 loc) · 3.03 KB
/
exp_scaletest_smtp.go
File metadata and controls
112 lines (99 loc) · 3.03 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
//go:build !slim
package cli
import (
"fmt"
"os/signal"
"time"
"golang.org/x/xerrors"
"cdr.dev/slog/v3"
"cdr.dev/slog/v3/sloggers/sloghuman"
"github.com/coder/coder/v2/scaletest/smtpmock"
"github.com/coder/serpent"
)
func (*RootCmd) scaletestSMTP() *serpent.Command {
var (
hostAddress string
smtpPort int64
apiPort int64
purgeAtCount int64
)
cmd := &serpent.Command{
Use: "smtp",
Short: "Start a mock SMTP server for testing",
Long: `Start a mock SMTP server with an HTTP API server that can be used to purge
messages and get messages by email.`,
Handler: func(inv *serpent.Invocation) error {
ctx := inv.Context()
notifyCtx, stop := signal.NotifyContext(ctx, StopSignals...)
defer stop()
ctx = notifyCtx
logger := slog.Make(sloghuman.Sink(inv.Stderr)).Leveled(slog.LevelInfo)
config := smtpmock.Config{
HostAddress: hostAddress,
SMTPPort: int(smtpPort),
APIPort: int(apiPort),
Logger: logger,
}
srv := new(smtpmock.Server)
if err := srv.Start(ctx, config); err != nil {
return xerrors.Errorf("start mock SMTP server: %w", err)
}
defer func() {
_ = srv.Stop()
}()
_, _ = fmt.Fprintf(inv.Stdout, "Mock SMTP server started on %s\n", srv.SMTPAddress())
_, _ = fmt.Fprintf(inv.Stdout, "HTTP API server started on %s\n", srv.APIAddress())
if purgeAtCount > 0 {
_, _ = fmt.Fprintf(inv.Stdout, " Auto-purge when message count reaches %d\n", purgeAtCount)
}
ticker := time.NewTicker(10 * time.Second)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
_, _ = fmt.Fprintf(inv.Stdout, "\nTotal messages received since last purge: %d\n", srv.MessageCount())
return nil
case <-ticker.C:
count := srv.MessageCount()
if count > 0 {
_, _ = fmt.Fprintf(inv.Stdout, "Messages received: %d\n", count)
}
if purgeAtCount > 0 && int64(count) >= purgeAtCount {
_, _ = fmt.Fprintf(inv.Stdout, "Message count (%d) reached threshold (%d). Purging...\n", count, purgeAtCount)
srv.Purge()
continue
}
}
}
},
}
cmd.Options = []serpent.Option{
{
Flag: "host-address",
Env: "CODER_SCALETEST_SMTP_HOST_ADDRESS",
Default: "localhost",
Description: "Host address to bind the mock SMTP and API servers.",
Value: serpent.StringOf(&hostAddress),
},
{
Flag: "smtp-port",
Env: "CODER_SCALETEST_SMTP_PORT",
Description: "Port for the mock SMTP server. Uses a random port if not specified.",
Value: serpent.Int64Of(&smtpPort),
},
{
Flag: "api-port",
Env: "CODER_SCALETEST_SMTP_API_PORT",
Description: "Port for the HTTP API server. Uses a random port if not specified.",
Value: serpent.Int64Of(&apiPort),
},
{
Flag: "purge-at-count",
Env: "CODER_SCALETEST_SMTP_PURGE_AT_COUNT",
Default: "100000",
Description: "Maximum number of messages to keep before auto-purging. Set to 0 to disable.",
Value: serpent.Int64Of(&purgeAtCount),
},
}
return cmd
}