| 1 | // SPDX-License-Identifier: GPL-3.0-or-later |
| 2 | |
| 3 | package jobmgr |
| 4 | |
| 5 | import ( |
| 6 | "bytes" |
| 7 | "context" |
| 8 | "errors" |
| 9 | "strings" |
| 10 | "sync" |
| 11 | "testing" |
| 12 | "time" |
| 13 | |
| 14 | "github.com/netdata/netdata/go/plugins/pkg/funcapi" |
| 15 | "github.com/netdata/netdata/go/plugins/pkg/netdataapi" |
| 16 | "github.com/netdata/netdata/go/plugins/pkg/safewriter" |
| 17 | "github.com/netdata/netdata/go/plugins/plugin/framework/collectorapi" |
| 18 | "github.com/netdata/netdata/go/plugins/plugin/framework/confgroup" |
| 19 | "github.com/netdata/netdata/go/plugins/plugin/framework/dyncfg" |
| 20 | |
| 21 | "github.com/stretchr/testify/assert" |
| 22 | "github.com/stretchr/testify/require" |
| 23 | ) |
| 24 | |
| 25 | const testPluginName = "test" |
| 26 | |
| 27 | type wantExposedEntry struct { |
| 28 | cfg confgroup.Config |
| 29 | status dyncfg.Status |
| 30 | } |
| 31 | |
| 32 | type runSim struct { |
| 33 | do func(mgr *Manager, in chan []*confgroup.Group) |
| 34 | |
| 35 | wantDiscovered []confgroup.Config |
| 36 | wantSeen []confgroup.Config |
| 37 | wantExposed []wantExposedEntry |
| 38 | wantRunning []string |
| 39 | wantDyncfg string |
| 40 | } |
| 41 | |
| 42 | const funcResultEndMarker = "FUNCTION_RESULT_END\n\n" |
| 43 | |
| 44 | type simOutput struct { |
| 45 | mu sync.Mutex |
| 46 | |
| 47 | buf bytes.Buffer |
| 48 | funcResultCount int |
| 49 | tail string |
| 50 | } |
| 51 | |
| 52 | func (o *simOutput) Write(p []byte) (int, error) { |
| 53 | o.mu.Lock() |
| 54 | defer o.mu.Unlock() |
| 55 | |
| 56 | n, err := o.buf.Write(p) |
| 57 | if n > 0 { |
| 58 | data := o.tail + string(p[:n]) |
| 59 | o.funcResultCount += strings.Count(data, funcResultEndMarker) |
| 60 | |
| 61 | tailLen := len(funcResultEndMarker) - 1 |
| 62 | if len(data) > tailLen { |
| 63 | o.tail = data[len(data)-tailLen:] |
| 64 | } else { |
| 65 | o.tail = data |
| 66 | } |
| 67 | } |
| 68 | |
| 69 | return n, err |
| 70 | } |
| 71 | |
| 72 | func (o *simOutput) String() string { |
| 73 | o.mu.Lock() |
| 74 | defer o.mu.Unlock() |
| 75 | return o.buf.String() |
| 76 | } |
| 77 | |
| 78 | func (o *simOutput) FuncResultCount() int { |
| 79 | o.mu.Lock() |
| 80 | defer o.mu.Unlock() |
| 81 | return o.funcResultCount |
| 82 | } |
| 83 | |
| 84 | func (s *runSim) run(t *testing.T) { |
| 85 | t.Helper() |
| 86 | |
| 87 | require.NotNil(t, s.do, "s.do is nil") |
| 88 | |
| 89 | var out simOutput |
| 90 | mgr := New(Config{PluginName: testPluginName}) |
| 91 | mgr.SetDyncfgResponder(dyncfg.NewResponder(netdataapi.New(safewriter.New(&out)))) |
| 92 | mgr.modules = prepareMockRegistry() |
| 93 | |
| 94 | done := make(chan struct{}) |
| 95 | grpCh := make(chan []*confgroup.Group) |
| 96 | ctx, cancel := context.WithCancel(context.Background()) |
| 97 | |
| 98 | go func() { defer close(done); defer close(grpCh); mgr.Run(ctx, grpCh) }() |
| 99 | |
| 100 | timeout := time.Second * 5 |
| 101 | |
| 102 | select { |
| 103 | case <-mgr.started: |
| 104 | case <-time.After(timeout): |
| 105 | t.Errorf("failed to start work in %s", timeout) |
| 106 | } |
| 107 | |
| 108 | s.do(mgr, grpCh) |
| 109 | |
| 110 | expectedResults := strings.Count(s.wantDyncfg, "FUNCTION_RESULT_END") |
| 111 | require.Eventually(t, func() bool { |
| 112 | return countDiscovered(mgr) == len(s.wantDiscovered) && |
| 113 | mgr.collectorSeen.Count() == len(s.wantSeen) && |
| 114 | mgr.collectorExposed.Count() == len(s.wantExposed) && |
| 115 | runningSetMatches(mgr.runningJobs.snapshot(), s.wantRunning) && |
| 116 | out.FuncResultCount() >= expectedResults |
| 117 | }, timeout, 10*time.Millisecond, "manager state did not settle before shutdown") |
| 118 | |
| 119 | runningBeforeShutdown := make(map[string]struct{}) |
| 120 | for _, job := range mgr.runningJobs.snapshot() { |
| 121 | runningBeforeShutdown[job.FullName()] = struct{}{} |
| 122 | } |
| 123 | |
| 124 | cancel() |
| 125 | |
| 126 | select { |
| 127 | case <-done: |
| 128 | case <-time.After(timeout): |
| 129 | t.Errorf("failed to finish work in %s", timeout) |
| 130 | } |
| 131 | |
| 132 | var lines []string |
| 133 | skipNextEmpty := false |
| 134 | for s := range strings.SplitSeq(out.String(), "\n") { |
| 135 | if strings.HasPrefix(s, "CONFIG") && strings.Contains(s, " template ") { |
| 136 | skipNextEmpty = false |
| 137 | continue |
| 138 | } |
| 139 | if strings.HasPrefix(s, "FUNCTION GLOBAL") { |
| 140 | skipNextEmpty = true |
| 141 | continue |
| 142 | } |
| 143 | if skipNextEmpty && s == "" { |
| 144 | skipNextEmpty = false |
| 145 | continue |
| 146 | } |
| 147 | skipNextEmpty = false |
| 148 | if strings.HasPrefix(s, "FUNCTION_RESULT_BEGIN") { |
| 149 | parts := strings.Fields(s) |
| 150 | s = strings.Join(parts[:len(parts)-1], " ") // remove timestamp |
| 151 | } |
| 152 | lines = append(lines, s) |
| 153 | } |
| 154 | wantDyncfg, gotDyncfg := strings.TrimSpace(s.wantDyncfg), strings.TrimSpace(strings.Join(lines, "\n")) |
| 155 | |
| 156 | //fmt.Println(gotDyncfg) |
| 157 | |
| 158 | assert.Equal(t, wantDyncfg, gotDyncfg, "dyncfg commands") |
| 159 | |
| 160 | var n int |
| 161 | for _, cfgs := range mgr.discoveredConfigs.items { |
| 162 | n += len(cfgs) |
| 163 | } |
| 164 | |
| 165 | wantLen, gotLen := len(s.wantDiscovered), n |
| 166 | require.Equalf(t, wantLen, gotLen, "discoveredConfigs: different len (want %d got %d)", wantLen, gotLen) |
| 167 | |
| 168 | for _, cfg := range s.wantDiscovered { |
| 169 | cfgs, ok := mgr.discoveredConfigs.items[cfg.Source()] |
| 170 | require.Truef(t, ok, "discoveredConfigs: source %s is not found", cfg.Source()) |
| 171 | _, ok = cfgs[cfg.Hash()] |
| 172 | require.Truef(t, ok, "discoveredConfigs: source %s config %d is not found", cfg.Source(), cfg.Hash()) |
| 173 | } |
| 174 | |
| 175 | wantLen, gotLen = len(s.wantSeen), mgr.collectorSeen.Count() |
| 176 | require.Equalf(t, wantLen, gotLen, "seen: different len (want %d got %d)", wantLen, gotLen) |
| 177 | |
| 178 | for _, cfg := range s.wantSeen { |
| 179 | _, ok := mgr.collectorSeen.Lookup(cfg) |
| 180 | require.Truef(t, ok, "seen: config '%s' is not found", cfg.UID()) |
| 181 | } |
| 182 | |
| 183 | wantLen, gotLen = len(s.wantExposed), mgr.collectorExposed.Count() |
| 184 | require.Equalf(t, wantLen, gotLen, "exposed: different len (want %d got %d)", wantLen, gotLen) |
| 185 | |
| 186 | for _, we := range s.wantExposed { |
| 187 | entry, ok := mgr.collectorExposed.LookupByKey(we.cfg.ExposedKey()) |
| 188 | require.Truef(t, ok && we.cfg.UID() == entry.Cfg.UID(), "exposed: config '%s' is not found", we.cfg.UID()) |
| 189 | require.Truef(t, we.status == entry.Status, "exposed: wrong status for '%s', want %s got %s", we.cfg.UID(), we.status, entry.Status) |
| 190 | } |
| 191 | |
| 192 | wantLen, gotLen = len(s.wantRunning), len(runningBeforeShutdown) |
| 193 | require.Equalf(t, wantLen, gotLen, "runningJobs: different len (want %d got %d)", wantLen, gotLen) |
| 194 | for _, name := range s.wantRunning { |
| 195 | _, ok := runningBeforeShutdown[name] |
| 196 | require.Truef(t, ok, "runningJobs: job '%s' is not found", name) |
| 197 | } |
| 198 | } |
| 199 | |
| 200 | func prepareMockRegistry() collectorapi.Registry { |
| 201 | reg := collectorapi.Registry{} |
| 202 | type config struct { |
| 203 | OptionOne string `yaml:"option_one" json:"option_one"` |
| 204 | OptionTwo int64 `yaml:"option_two" json:"option_two"` |
| 205 | } |
| 206 | |
| 207 | reg.Register("success", collectorapi.Creator{ |
| 208 | JobConfigSchema: collectorapi.MockConfigSchema, |
| 209 | Create: func() collectorapi.CollectorV1 { |
| 210 | return &collectorapi.MockCollectorV1{ |
| 211 | ChartsFunc: func() *collectorapi.Charts { |
| 212 | return &collectorapi.Charts{&collectorapi.Chart{ID: "id", Title: "title", Units: "units", Dims: collectorapi.Dims{{ID: "id1"}}}} |
| 213 | }, |
| 214 | CollectFunc: func(context.Context) map[string]int64 { return map[string]int64{"id1": 1} }, |
| 215 | } |
| 216 | }, |
| 217 | Config: func() any { |
| 218 | return &config{OptionOne: "one", OptionTwo: 2} |
| 219 | }, |
| 220 | }) |
| 221 | reg.Register("fail", collectorapi.Creator{ |
| 222 | Create: func() collectorapi.CollectorV1 { |
| 223 | return &collectorapi.MockCollectorV1{ |
| 224 | InitFunc: func(context.Context) error { return errors.New("mock failed init") }, |
| 225 | } |
| 226 | }, |
| 227 | }) |
| 228 | |
| 229 | // CollectorV1 without Methods - for testing function_only config rejection |
| 230 | reg.Register("nofuncs", collectorapi.Creator{ |
| 231 | Create: func() collectorapi.CollectorV1 { |
| 232 | return &collectorapi.MockCollectorV1{ |
| 233 | ChartsFunc: func() *collectorapi.Charts { |
| 234 | return &collectorapi.Charts{&collectorapi.Chart{ID: "id", Title: "title", Units: "units", Dims: collectorapi.Dims{{ID: "id1"}}}} |
| 235 | }, |
| 236 | CollectFunc: func(context.Context) map[string]int64 { return map[string]int64{"id1": 1} }, |
| 237 | } |
| 238 | }, |
| 239 | }) |
| 240 | |
| 241 | // CollectorV1 with Methods - for testing config-level function_only |
| 242 | reg.Register("withfuncs", collectorapi.Creator{ |
| 243 | Create: func() collectorapi.CollectorV1 { |
| 244 | return &collectorapi.MockCollectorV1{ |
| 245 | ChartsFunc: func() *collectorapi.Charts { |
| 246 | return &collectorapi.Charts{&collectorapi.Chart{ID: "id", Title: "title", Units: "units", Dims: collectorapi.Dims{{ID: "id1"}}}} |
| 247 | }, |
| 248 | CollectFunc: func(context.Context) map[string]int64 { return map[string]int64{"id1": 1} }, |
| 249 | } |
| 250 | }, |
| 251 | Methods: func() []funcapi.MethodConfig { |
| 252 | return []funcapi.MethodConfig{{ID: "test-method", Name: "Test Method"}} |
| 253 | }, |
| 254 | }) |
| 255 | |
| 256 | // FunctionOnly module - for testing module-level function-only |
| 257 | reg.Register("funconly", collectorapi.Creator{ |
| 258 | FunctionOnly: true, |
| 259 | Create: func() collectorapi.CollectorV1 { |
| 260 | return &collectorapi.MockCollectorV1{ |
| 261 | ChartsFunc: func() *collectorapi.Charts { return nil }, |
| 262 | CollectFunc: func(context.Context) map[string]int64 { return nil }, |
| 263 | } |
| 264 | }, |
| 265 | Methods: func() []funcapi.MethodConfig { |
| 266 | return []funcapi.MethodConfig{{ID: "test-method", Name: "Test Method"}} |
| 267 | }, |
| 268 | }) |
| 269 | |
| 270 | return reg |
| 271 | } |
| 272 | |
| 273 | func countDiscovered(mgr *Manager) int { |
| 274 | var n int |
| 275 | for _, cfgs := range mgr.discoveredConfigs.items { |
| 276 | n += len(cfgs) |
| 277 | } |
| 278 | return n |
| 279 | } |
| 280 | |
| 281 | func runningSetMatches(jobs []runtimeJob, want []string) bool { |
| 282 | if len(jobs) != len(want) { |
| 283 | return false |
| 284 | } |
| 285 | |
| 286 | wantSet := make(map[string]struct{}, len(want)) |
| 287 | for _, name := range want { |
| 288 | wantSet[name] = struct{}{} |
| 289 | } |
| 290 | |
| 291 | for _, job := range jobs { |
| 292 | if _, ok := wantSet[job.FullName()]; !ok { |
| 293 | return false |
| 294 | } |
| 295 | } |
| 296 | |
| 297 | return true |
| 298 | } |