| 1 | // SPDX-License-Identifier: GPL-3.0-or-later |
| 2 | |
| 3 | package jobmgr |
| 4 | |
| 5 | import ( |
| 6 | "bytes" |
| 7 | "context" |
| 8 | "encoding/json" |
| 9 | "sync" |
| 10 | "testing" |
| 11 | "time" |
| 12 | |
| 13 | "github.com/stretchr/testify/assert" |
| 14 | "github.com/stretchr/testify/require" |
| 15 | |
| 16 | "github.com/netdata/netdata/go/plugins/pkg/funcapi" |
| 17 | "github.com/netdata/netdata/go/plugins/pkg/netdataapi" |
| 18 | "github.com/netdata/netdata/go/plugins/pkg/safewriter" |
| 19 | "github.com/netdata/netdata/go/plugins/plugin/framework/collectorapi" |
| 20 | "github.com/netdata/netdata/go/plugins/plugin/framework/dyncfg" |
| 21 | "github.com/netdata/netdata/go/plugins/plugin/framework/functions" |
| 22 | ) |
| 23 | |
| 24 | type mockMethodHandler struct { |
| 25 | paramsFunc func(ctx context.Context, method string) ([]funcapi.ParamConfig, error) |
| 26 | handleFunc func(ctx context.Context, method string, params funcapi.ResolvedParams) *funcapi.FunctionResponse |
| 27 | } |
| 28 | |
| 29 | func (m *mockMethodHandler) MethodParams(ctx context.Context, method string) ([]funcapi.ParamConfig, error) { |
| 30 | if m.paramsFunc != nil { |
| 31 | return m.paramsFunc(ctx, method) |
| 32 | } |
| 33 | return nil, nil |
| 34 | } |
| 35 | |
| 36 | func (m *mockMethodHandler) Handle(ctx context.Context, method string, params funcapi.ResolvedParams) *funcapi.FunctionResponse { |
| 37 | if m.handleFunc != nil { |
| 38 | return m.handleFunc(ctx, method, params) |
| 39 | } |
| 40 | return nil |
| 41 | } |
| 42 | |
| 43 | func (m *mockMethodHandler) Cleanup(context.Context) {} |
| 44 | |
| 45 | func TestExecuteFunction_ModuleMethodPaths(t *testing.T) { |
| 46 | tests := map[string]struct { |
| 47 | methods []funcapi.MethodConfig |
| 48 | fnArgs []string |
| 49 | fnPayload map[string]any |
| 50 | generationRace bool |
| 51 | wantStatus int |
| 52 | wantHelp string |
| 53 | wantAccepted []string |
| 54 | wantRequiredIDs []string |
| 55 | wantDataValue any |
| 56 | wantResolvedJob string |
| 57 | wantErrorContains string |
| 58 | }{ |
| 59 | "success with __job resolution": { |
| 60 | methods: []funcapi.MethodConfig{{ID: "details", Help: "details help"}}, |
| 61 | fnArgs: []string{"__job:job1"}, |
| 62 | wantStatus: 200, |
| 63 | wantHelp: "details help", |
| 64 | wantAccepted: []string{"__job"}, |
| 65 | wantRequiredIDs: []string{"__job"}, |
| 66 | wantDataValue: "row", |
| 67 | wantResolvedJob: "job1", |
| 68 | }, |
| 69 | "info response includes module and method params": { |
| 70 | methods: []funcapi.MethodConfig{{ |
| 71 | ID: "details", |
| 72 | Help: "details help", |
| 73 | RequiredParams: []funcapi.ParamConfig{{ |
| 74 | ID: "scope", |
| 75 | Name: "Scope", |
| 76 | Selection: funcapi.ParamSelect, |
| 77 | Options: []funcapi.ParamOption{{ID: "default", Name: "Default"}}, |
| 78 | }}, |
| 79 | }}, |
| 80 | fnArgs: []string{"info"}, |
| 81 | wantStatus: 200, |
| 82 | wantHelp: "details help", |
| 83 | wantAccepted: []string{"__job", "scope"}, |
| 84 | wantRequiredIDs: []string{"__job", "scope"}, |
| 85 | }, |
| 86 | "generation race returns 503": { |
| 87 | methods: []funcapi.MethodConfig{{ID: "details"}}, |
| 88 | fnArgs: []string{"__job:job1"}, |
| 89 | generationRace: true, |
| 90 | wantStatus: 503, |
| 91 | wantErrorContains: "replaced during request", |
| 92 | }, |
| 93 | "explicit unknown __job in args returns 404": { |
| 94 | methods: []funcapi.MethodConfig{{ID: "details"}}, |
| 95 | fnArgs: []string{"__job:missing"}, |
| 96 | wantStatus: 404, |
| 97 | wantErrorContains: "unknown job 'missing'", |
| 98 | }, |
| 99 | "explicit unknown __job in payload returns 404": { |
| 100 | methods: []funcapi.MethodConfig{{ID: "details"}}, |
| 101 | fnPayload: map[string]any{"__job": "missing"}, |
| 102 | wantStatus: 404, |
| 103 | wantErrorContains: "unknown job 'missing'", |
| 104 | }, |
| 105 | "multiple __job values return 400": { |
| 106 | methods: []funcapi.MethodConfig{{ID: "details"}}, |
| 107 | fnPayload: map[string]any{"__job": []string{"job1", "job2"}}, |
| 108 | wantStatus: 400, |
| 109 | wantErrorContains: "parameter '__job' expects a single value", |
| 110 | }, |
| 111 | } |
| 112 | |
| 113 | for name, tc := range tests { |
| 114 | t.Run(name, func(t *testing.T) { |
| 115 | writer := &jsonWriteCapture{} |
| 116 | var gotJob string |
| 117 | var mgr *Manager |
| 118 | |
| 119 | methodHandler := &mockMethodHandler{} |
| 120 | if len(tc.fnArgs) == 0 || tc.fnArgs[0] != "info" { |
| 121 | methodHandler.handleFunc = func(ctx context.Context, method string, params funcapi.ResolvedParams) *funcapi.FunctionResponse { |
| 122 | if tc.generationRace { |
| 123 | mgr.funcCtl.OnJobStart(&lockProbeJob{fullName: "mod_job1", moduleName: "mod", name: "job1"}) |
| 124 | return &funcapi.FunctionResponse{Status: 200, Help: "should be replaced"} |
| 125 | } |
| 126 | |
| 127 | gotJob = params.GetOne("__job") |
| 128 | return &funcapi.FunctionResponse{ |
| 129 | Status: 200, |
| 130 | Help: tc.wantHelp, |
| 131 | Columns: map[string]any{ |
| 132 | "value": map[string]any{"name": "Value"}, |
| 133 | }, |
| 134 | Data: [][]any{{tc.wantDataValue}}, |
| 135 | } |
| 136 | } |
| 137 | } |
| 138 | |
| 139 | mgr = newModuleDispatchTestManager(t, nil, writer.write, methodHandler, tc.methods) |
| 140 | var payload []byte |
| 141 | if tc.fnPayload != nil { |
| 142 | var err error |
| 143 | payload, err = json.Marshal(tc.fnPayload) |
| 144 | require.NoError(t, err) |
| 145 | } |
| 146 | mgr.ExecuteFunction("mod:details", functions.Function{ |
| 147 | UID: "module-test", |
| 148 | Timeout: time.Second, |
| 149 | Args: tc.fnArgs, |
| 150 | Payload: payload, |
| 151 | }) |
| 152 | |
| 153 | resp := writer.requireResponse(t) |
| 154 | assert.Equal(t, tc.wantStatus, writer.code) |
| 155 | assert.Equal(t, float64(tc.wantStatus), resp["status"]) |
| 156 | if tc.wantErrorContains != "" { |
| 157 | assert.Contains(t, resp["errorMessage"], tc.wantErrorContains) |
| 158 | return |
| 159 | } |
| 160 | |
| 161 | if tc.wantResolvedJob != "" { |
| 162 | assert.Equal(t, tc.wantResolvedJob, gotJob) |
| 163 | } |
| 164 | if tc.wantHelp != "" { |
| 165 | assert.Equal(t, tc.wantHelp, resp["help"]) |
| 166 | } |
| 167 | assert.Equal(t, tc.wantAccepted, jsonArrayStrings(t, resp["accepted_params"])) |
| 168 | |
| 169 | required := jsonObjectArray(t, resp["required_params"]) |
| 170 | require.Len(t, required, len(tc.wantRequiredIDs)) |
| 171 | for i, id := range tc.wantRequiredIDs { |
| 172 | assert.Equal(t, id, required[i]["id"]) |
| 173 | } |
| 174 | if tc.wantDataValue != nil { |
| 175 | assert.Equal(t, tc.wantDataValue, jsonNestedArrayValue(t, resp["data"], 0, 0)) |
| 176 | } |
| 177 | }) |
| 178 | } |
| 179 | } |
| 180 | |
| 181 | func TestExecuteFunction_ContextBehavior(t *testing.T) { |
| 182 | tests := map[string]struct { |
| 183 | managerCtx context.Context |
| 184 | wantMarker string |
| 185 | wantMarkerSet bool |
| 186 | wantHasDeadline bool |
| 187 | }{ |
| 188 | "uses background fallback before manager context is set": { |
| 189 | wantHasDeadline: true, |
| 190 | }, |
| 191 | "uses manager context when available": { |
| 192 | managerCtx: context.WithValue(context.Background(), dispatchContextKey("marker"), "manager"), |
| 193 | wantMarker: "manager", |
| 194 | wantMarkerSet: true, |
| 195 | wantHasDeadline: true, |
| 196 | }, |
| 197 | } |
| 198 | |
| 199 | for name, tc := range tests { |
| 200 | t.Run(name, func(t *testing.T) { |
| 201 | writer := &jsonWriteCapture{} |
| 202 | mgr := newModuleDispatchTestManager(t, nil, writer.write, &mockMethodHandler{ |
| 203 | handleFunc: func(ctx context.Context, method string, params funcapi.ResolvedParams) *funcapi.FunctionResponse { |
| 204 | require.NotNil(t, ctx) |
| 205 | _, hasDeadline := ctx.Deadline() |
| 206 | assert.Equal(t, tc.wantHasDeadline, hasDeadline) |
| 207 | gotMarker := ctx.Value(dispatchContextKey("marker")) |
| 208 | if tc.wantMarkerSet { |
| 209 | assert.Equal(t, tc.wantMarker, gotMarker) |
| 210 | } else { |
| 211 | assert.Nil(t, gotMarker) |
| 212 | } |
| 213 | return &funcapi.FunctionResponse{Status: 200} |
| 214 | }, |
| 215 | }, []funcapi.MethodConfig{{ID: "details"}}) |
| 216 | if tc.managerCtx != nil { |
| 217 | mgr.funcCtl.Init(tc.managerCtx) |
| 218 | } |
| 219 | |
| 220 | mgr.ExecuteFunction("mod:details", functions.Function{ |
| 221 | UID: "module-context", |
| 222 | Timeout: time.Second, |
| 223 | Args: []string{"__job:job1"}, |
| 224 | }) |
| 225 | |
| 226 | resp := writer.requireResponse(t) |
| 227 | assert.Equal(t, float64(200), resp["status"]) |
| 228 | }) |
| 229 | } |
| 230 | } |
| 231 | |
| 232 | func TestJobMethodRegisteredHandlerPaths(t *testing.T) { |
| 233 | tests := map[string]struct { |
| 234 | fnArgs []string |
| 235 | wantRequiredLen int |
| 236 | wantDataValue any |
| 237 | }{ |
| 238 | "success path omits __job": { |
| 239 | wantRequiredLen: 0, |
| 240 | wantDataValue: "job1", |
| 241 | }, |
| 242 | "info path omits __job": { |
| 243 | fnArgs: []string{"info"}, |
| 244 | wantRequiredLen: 0, |
| 245 | }, |
| 246 | } |
| 247 | |
| 248 | for name, tc := range tests { |
| 249 | t.Run(name, func(t *testing.T) { |
| 250 | writer := &jsonWriteCapture{} |
| 251 | fnReg := newCapturingFunctionRegistry() |
| 252 | mgr := newJobMethodDispatchTestManager(t, fnReg, writer.write, &mockMethodHandler{ |
| 253 | handleFunc: func(ctx context.Context, method string, params funcapi.ResolvedParams) *funcapi.FunctionResponse { |
| 254 | return &funcapi.FunctionResponse{ |
| 255 | Status: 200, |
| 256 | Help: "job details help", |
| 257 | Columns: map[string]any{ |
| 258 | "value": map[string]any{"name": "Value"}, |
| 259 | }, |
| 260 | Data: [][]any{{"job1"}}, |
| 261 | } |
| 262 | }, |
| 263 | }, []funcapi.MethodConfig{{ID: "job-details", Help: "job details help"}}) |
| 264 | |
| 265 | handler := fnReg.requireHandler(t, "mod:job-details") |
| 266 | handler(functions.Function{ |
| 267 | UID: "job-handler", |
| 268 | Timeout: time.Second, |
| 269 | Args: tc.fnArgs, |
| 270 | }) |
| 271 | |
| 272 | resp := writer.requireResponse(t) |
| 273 | assert.Equal(t, float64(200), resp["status"]) |
| 274 | assert.NotContains(t, jsonArrayStrings(t, resp["accepted_params"]), "__job") |
| 275 | assert.Len(t, jsonObjectArray(t, resp["required_params"]), tc.wantRequiredLen) |
| 276 | if tc.wantDataValue != nil { |
| 277 | assert.Equal(t, tc.wantDataValue, jsonNestedArrayValue(t, resp["data"], 0, 0)) |
| 278 | } |
| 279 | mgr.stopRunningJob("mod_job1") |
| 280 | }) |
| 281 | } |
| 282 | } |
| 283 | |
| 284 | func TestFunctionDispatch_ResponsePaths(t *testing.T) { |
| 285 | tests := map[string]struct { |
| 286 | useJSONWriter bool |
| 287 | rebindResponder bool |
| 288 | nilRebindResponder bool |
| 289 | marshalFail bool |
| 290 | wantWriterStatus int |
| 291 | wantResponderUID string |
| 292 | wantResponderJSON string |
| 293 | wantFirstUID string |
| 294 | wantSecondUID string |
| 295 | }{ |
| 296 | "JSONWriter takes precedence when configured": { |
| 297 | useJSONWriter: true, |
| 298 | wantWriterStatus: 200, |
| 299 | wantResponderUID: "writer-first", |
| 300 | }, |
| 301 | "responder fallback is used when JSONWriter is nil": { |
| 302 | wantResponderUID: "responder-fallback", |
| 303 | wantResponderJSON: "\"status\":200", |
| 304 | }, |
| 305 | "responder rebinding updates only the responder-backed path": { |
| 306 | rebindResponder: true, |
| 307 | wantFirstUID: "before-rebind", |
| 308 | wantSecondUID: "after-rebind", |
| 309 | }, |
| 310 | "nil responder rebinding preserves the current responder-backed path": { |
| 311 | nilRebindResponder: true, |
| 312 | wantFirstUID: "before-nil-rebind", |
| 313 | wantSecondUID: "after-nil-rebind", |
| 314 | }, |
| 315 | "marshal failure falls back to JSONWriter with 500": { |
| 316 | useJSONWriter: true, |
| 317 | marshalFail: true, |
| 318 | wantWriterStatus: 500, |
| 319 | wantResponderUID: "writer-marshal-fail", |
| 320 | }, |
| 321 | "marshal failure falls back to responder with 500": { |
| 322 | marshalFail: true, |
| 323 | wantResponderUID: "responder-marshal-fail", |
| 324 | wantResponderJSON: "\"status\":500", |
| 325 | }, |
| 326 | } |
| 327 | |
| 328 | for name, tc := range tests { |
| 329 | t.Run(name, func(t *testing.T) { |
| 330 | handler := &mockMethodHandler{ |
| 331 | handleFunc: func(ctx context.Context, method string, params funcapi.ResolvedParams) *funcapi.FunctionResponse { |
| 332 | if tc.marshalFail { |
| 333 | return &funcapi.FunctionResponse{ |
| 334 | Status: 200, |
| 335 | Data: [][]any{{make(chan int)}}, |
| 336 | } |
| 337 | } |
| 338 | return &funcapi.FunctionResponse{Status: 200} |
| 339 | }, |
| 340 | } |
| 341 | |
| 342 | if tc.useJSONWriter { |
| 343 | writer := &jsonWriteCapture{} |
| 344 | var responderOut bytes.Buffer |
| 345 | |
| 346 | mgr := newModuleDispatchTestManager(t, dyncfg.NewResponder(netdataapi.New(safewriter.New(&responderOut))), writer.write, handler, []funcapi.MethodConfig{{ID: "details"}}) |
| 347 | mgr.ExecuteFunction("mod:details", functions.Function{ |
| 348 | UID: tc.wantResponderUID, |
| 349 | Timeout: time.Second, |
| 350 | Args: []string{"__job:job1"}, |
| 351 | }) |
| 352 | |
| 353 | resp := writer.requireResponse(t) |
| 354 | assert.Equal(t, float64(tc.wantWriterStatus), resp["status"]) |
| 355 | if tc.marshalFail { |
| 356 | assert.Contains(t, resp["errorMessage"], "failed to encode response") |
| 357 | } |
| 358 | assert.NotContains(t, responderOut.String(), "FUNCTION_RESULT_BEGIN "+tc.wantResponderUID) |
| 359 | return |
| 360 | } |
| 361 | |
| 362 | if tc.rebindResponder { |
| 363 | var firstOut bytes.Buffer |
| 364 | var secondOut bytes.Buffer |
| 365 | |
| 366 | mgr := newModuleDispatchTestManager(t, dyncfg.NewResponder(netdataapi.New(safewriter.New(&firstOut))), nil, handler, []funcapi.MethodConfig{{ID: "details"}}) |
| 367 | mgr.ExecuteFunction("mod:details", functions.Function{ |
| 368 | UID: tc.wantFirstUID, |
| 369 | Timeout: time.Second, |
| 370 | Args: []string{"__job:job1"}, |
| 371 | }) |
| 372 | |
| 373 | mgr.SetDyncfgResponder(dyncfg.NewResponder(netdataapi.New(safewriter.New(&secondOut)))) |
| 374 | mgr.ExecuteFunction("mod:details", functions.Function{ |
| 375 | UID: tc.wantSecondUID, |
| 376 | Timeout: time.Second, |
| 377 | Args: []string{"__job:job1"}, |
| 378 | }) |
| 379 | |
| 380 | assert.Contains(t, firstOut.String(), "FUNCTION_RESULT_BEGIN "+tc.wantFirstUID) |
| 381 | assert.NotContains(t, firstOut.String(), "FUNCTION_RESULT_BEGIN "+tc.wantSecondUID) |
| 382 | assert.Contains(t, secondOut.String(), "FUNCTION_RESULT_BEGIN "+tc.wantSecondUID) |
| 383 | return |
| 384 | } |
| 385 | |
| 386 | if tc.nilRebindResponder { |
| 387 | var responderOut bytes.Buffer |
| 388 | |
| 389 | mgr := newModuleDispatchTestManager(t, dyncfg.NewResponder(netdataapi.New(safewriter.New(&responderOut))), nil, handler, []funcapi.MethodConfig{{ID: "details"}}) |
| 390 | mgr.ExecuteFunction("mod:details", functions.Function{ |
| 391 | UID: tc.wantFirstUID, |
| 392 | Timeout: time.Second, |
| 393 | Args: []string{"__job:job1"}, |
| 394 | }) |
| 395 | |
| 396 | mgr.SetDyncfgResponder(nil) |
| 397 | mgr.ExecuteFunction("mod:details", functions.Function{ |
| 398 | UID: tc.wantSecondUID, |
| 399 | Timeout: time.Second, |
| 400 | Args: []string{"__job:job1"}, |
| 401 | }) |
| 402 | |
| 403 | assert.Contains(t, responderOut.String(), "FUNCTION_RESULT_BEGIN "+tc.wantFirstUID) |
| 404 | assert.Contains(t, responderOut.String(), "FUNCTION_RESULT_BEGIN "+tc.wantSecondUID) |
| 405 | return |
| 406 | } |
| 407 | |
| 408 | var responderOut bytes.Buffer |
| 409 | mgr := newModuleDispatchTestManager(t, dyncfg.NewResponder(netdataapi.New(safewriter.New(&responderOut))), nil, handler, []funcapi.MethodConfig{{ID: "details"}}) |
| 410 | mgr.ExecuteFunction("mod:details", functions.Function{ |
| 411 | UID: tc.wantResponderUID, |
| 412 | Timeout: time.Second, |
| 413 | Args: []string{"__job:job1"}, |
| 414 | }) |
| 415 | |
| 416 | assert.Contains(t, responderOut.String(), "FUNCTION_RESULT_BEGIN "+tc.wantResponderUID) |
| 417 | assert.Contains(t, responderOut.String(), tc.wantResponderJSON) |
| 418 | if tc.marshalFail { |
| 419 | assert.Contains(t, responderOut.String(), "failed to encode response") |
| 420 | } |
| 421 | }) |
| 422 | } |
| 423 | } |
| 424 | |
| 425 | func TestCleanup_UnregistersStaticFunctionsBeforeStoppingJobs(t *testing.T) { |
| 426 | fnReg := newCapturingFunctionRegistry() |
| 427 | mgr := New(Config{PluginName: testPluginName, FnReg: fnReg}) |
| 428 | |
| 429 | staticCreator := collectorapi.Creator{ |
| 430 | Methods: func() []funcapi.MethodConfig { |
| 431 | return []funcapi.MethodConfig{{ID: "static-method"}} |
| 432 | }, |
| 433 | } |
| 434 | jobCreator := collectorapi.Creator{ |
| 435 | JobMethods: func(_ collectorapi.RuntimeJob) []funcapi.MethodConfig { |
| 436 | return []funcapi.MethodConfig{{ID: "job-method"}} |
| 437 | }, |
| 438 | } |
| 439 | |
| 440 | mgr.modules = collectorapi.Registry{ |
| 441 | "staticmod": staticCreator, |
| 442 | "jobmod": jobCreator, |
| 443 | } |
| 444 | mgr.funcCtl.RegisterModules(mgr.modules) |
| 445 | |
| 446 | mgr.startRunningJob(&lockProbeJob{fullName: "staticmod_job1", moduleName: "staticmod", name: "job1"}) |
| 447 | mgr.startRunningJob(&lockProbeJob{fullName: "jobmod_job1", moduleName: "jobmod", name: "job1"}) |
| 448 | |
| 449 | mgr.cleanup() |
| 450 | |
| 451 | unregistered := fnReg.unregisteredNames() |
| 452 | assert.Contains(t, unregistered, "staticmod:static-method") |
| 453 | assert.Contains(t, unregistered, "jobmod:job-method") |
| 454 | assert.Less( |
| 455 | t, |
| 456 | fnReg.unregisteredIndex("staticmod:static-method"), |
| 457 | fnReg.unregisteredIndex("jobmod:job-method"), |
| 458 | "static module cleanup must run before per-job stop cleanup", |
| 459 | ) |
| 460 | } |
| 461 | |
| 462 | type dispatchContextKey string |
| 463 | |
| 464 | type jsonWriteCapture struct { |
| 465 | calls int |
| 466 | code int |
| 467 | raw []byte |
| 468 | } |
| 469 | |
| 470 | func (c *jsonWriteCapture) write(payload []byte, code int) { |
| 471 | c.calls++ |
| 472 | c.code = code |
| 473 | c.raw = append([]byte(nil), payload...) |
| 474 | } |
| 475 | |
| 476 | func (c *jsonWriteCapture) requireResponse(t *testing.T) map[string]any { |
| 477 | t.Helper() |
| 478 | require.Equal(t, 1, c.calls, "expected exactly one JSON writer call") |
| 479 | |
| 480 | var resp map[string]any |
| 481 | require.NoError(t, json.Unmarshal(c.raw, &resp)) |
| 482 | return resp |
| 483 | } |
| 484 | |
| 485 | type capturingFunctionRegistry struct { |
| 486 | mu sync.Mutex |
| 487 | handlers map[string]func(functions.Function) |
| 488 | unregistered []string |
| 489 | } |
| 490 | |
| 491 | func newCapturingFunctionRegistry() *capturingFunctionRegistry { |
| 492 | return &capturingFunctionRegistry{ |
| 493 | handlers: make(map[string]func(functions.Function)), |
| 494 | } |
| 495 | } |
| 496 | |
| 497 | func (r *capturingFunctionRegistry) Register(name string, fn func(functions.Function)) { |
| 498 | r.mu.Lock() |
| 499 | r.handlers[name] = fn |
| 500 | r.mu.Unlock() |
| 501 | } |
| 502 | |
| 503 | func (r *capturingFunctionRegistry) Unregister(name string) { |
| 504 | r.mu.Lock() |
| 505 | r.unregistered = append(r.unregistered, name) |
| 506 | delete(r.handlers, name) |
| 507 | r.mu.Unlock() |
| 508 | } |
| 509 | |
| 510 | func (r *capturingFunctionRegistry) RegisterPrefix(string, string, func(functions.Function)) {} |
| 511 | func (r *capturingFunctionRegistry) UnregisterPrefix(string, string) {} |
| 512 | |
| 513 | func (r *capturingFunctionRegistry) requireHandler(t *testing.T, name string) func(functions.Function) { |
| 514 | t.Helper() |
| 515 | r.mu.Lock() |
| 516 | defer r.mu.Unlock() |
| 517 | |
| 518 | handler, ok := r.handlers[name] |
| 519 | require.True(t, ok, "handler %q was not registered", name) |
| 520 | return handler |
| 521 | } |
| 522 | |
| 523 | func (r *capturingFunctionRegistry) unregisteredNames() []string { |
| 524 | r.mu.Lock() |
| 525 | defer r.mu.Unlock() |
| 526 | |
| 527 | out := make([]string, len(r.unregistered)) |
| 528 | copy(out, r.unregistered) |
| 529 | return out |
| 530 | } |
| 531 | |
| 532 | func (r *capturingFunctionRegistry) unregisteredIndex(name string) int { |
| 533 | r.mu.Lock() |
| 534 | defer r.mu.Unlock() |
| 535 | |
| 536 | for i, got := range r.unregistered { |
| 537 | if got == name { |
| 538 | return i |
| 539 | } |
| 540 | } |
| 541 | return -1 |
| 542 | } |
| 543 | |
| 544 | func newModuleDispatchTestManager( |
| 545 | t *testing.T, |
| 546 | api *dyncfg.Responder, |
| 547 | jsonWriter func([]byte, int), |
| 548 | methodHandler funcapi.MethodHandler, |
| 549 | methods []funcapi.MethodConfig, |
| 550 | ) *Manager { |
| 551 | t.Helper() |
| 552 | |
| 553 | mgr := New(Config{ |
| 554 | PluginName: testPluginName, |
| 555 | FunctionJSONWriter: jsonWriter, |
| 556 | }) |
| 557 | if api != nil { |
| 558 | mgr.SetDyncfgResponder(api) |
| 559 | } |
| 560 | |
| 561 | creator := collectorapi.Creator{ |
| 562 | Methods: func() []funcapi.MethodConfig { return methods }, |
| 563 | MethodHandler: func(job collectorapi.RuntimeJob) funcapi.MethodHandler { |
| 564 | return methodHandler |
| 565 | }, |
| 566 | } |
| 567 | mgr.modules = collectorapi.Registry{"mod": creator} |
| 568 | mgr.funcCtl.RegisterModules(mgr.modules) |
| 569 | mgr.funcCtl.OnJobStart(&lockProbeJob{fullName: "mod_job1", moduleName: "mod", name: "job1"}) |
| 570 | |
| 571 | return mgr |
| 572 | } |
| 573 | |
| 574 | func newJobMethodDispatchTestManager( |
| 575 | t *testing.T, |
| 576 | fnReg FunctionRegistry, |
| 577 | jsonWriter func([]byte, int), |
| 578 | methodHandler funcapi.MethodHandler, |
| 579 | methods []funcapi.MethodConfig, |
| 580 | ) *Manager { |
| 581 | t.Helper() |
| 582 | |
| 583 | mgr := New(Config{ |
| 584 | PluginName: testPluginName, |
| 585 | FnReg: fnReg, |
| 586 | FunctionJSONWriter: jsonWriter, |
| 587 | }) |
| 588 | |
| 589 | creator := collectorapi.Creator{ |
| 590 | JobMethods: func(_ collectorapi.RuntimeJob) []funcapi.MethodConfig { return methods }, |
| 591 | MethodHandler: func(job collectorapi.RuntimeJob) funcapi.MethodHandler { |
| 592 | return methodHandler |
| 593 | }, |
| 594 | } |
| 595 | mgr.modules = collectorapi.Registry{"mod": creator} |
| 596 | mgr.funcCtl.RegisterModules(mgr.modules) |
| 597 | mgr.startRunningJob(&lockProbeJob{fullName: "mod_job1", moduleName: "mod", name: "job1"}) |
| 598 | |
| 599 | return mgr |
| 600 | } |
| 601 | |
| 602 | func jsonArrayStrings(t *testing.T, raw any) []string { |
| 603 | t.Helper() |
| 604 | |
| 605 | items, ok := raw.([]any) |
| 606 | require.True(t, ok, "expected []any, got %T", raw) |
| 607 | |
| 608 | out := make([]string, 0, len(items)) |
| 609 | for _, item := range items { |
| 610 | s, ok := item.(string) |
| 611 | require.True(t, ok, "expected string item, got %T", item) |
| 612 | out = append(out, s) |
| 613 | } |
| 614 | return out |
| 615 | } |
| 616 | |
| 617 | func jsonObjectArray(t *testing.T, raw any) []map[string]any { |
| 618 | t.Helper() |
| 619 | |
| 620 | items, ok := raw.([]any) |
| 621 | require.True(t, ok, "expected []any, got %T", raw) |
| 622 | |
| 623 | out := make([]map[string]any, 0, len(items)) |
| 624 | for _, item := range items { |
| 625 | obj, ok := item.(map[string]any) |
| 626 | require.True(t, ok, "expected map[string]any item, got %T", item) |
| 627 | out = append(out, obj) |
| 628 | } |
| 629 | return out |
| 630 | } |
| 631 | |
| 632 | func jsonNestedArrayValue(t *testing.T, raw any, row, col int) any { |
| 633 | t.Helper() |
| 634 | |
| 635 | rows, ok := raw.([]any) |
| 636 | require.True(t, ok, "expected []any rows, got %T", raw) |
| 637 | require.Len(t, rows, row+1) |
| 638 | |
| 639 | cols, ok := rows[row].([]any) |
| 640 | require.True(t, ok, "expected []any columns, got %T", rows[row]) |
| 641 | require.Len(t, cols, col+1) |
| 642 | |
| 643 | return cols[col] |
| 644 | } |