diff --git a/exporter/trace.go b/exporter/trace.go index 9e74852..ab06eb4 100644 --- a/exporter/trace.go +++ b/exporter/trace.go @@ -8,7 +8,6 @@ import ( "encoding/json" "errors" "fmt" - "io/fs" "log" "net/http" "os" @@ -109,22 +108,11 @@ type TraceCollector struct { func NewTraceCollector(config *Config) *TraceCollector { traceConfig := config.TraceConf - templateRootDir := "." - // path to look for the /templates directory. Defaults to cwd - if path, ok := os.LookupEnv("TRACE_ROOT_PATH"); ok { - templateRootDir = path - } - traceDir := "" - err := filepath.WalkDir(templateRootDir, func(path string, d fs.DirEntry, err error) error { - if err == nil && d.IsDir() && d.Name() == templateDirName { - traceDir = path - return nil - } - return nil - }) - if err != nil || traceDir == "" { + traceDir := detectTraceTemplatePath() + if traceDir == "" { log.Fatal("no template found") } + slog.Debug("using trace template path: " + traceDir) return &TraceCollector{ ProcessFetcher: NewAtomicProFetcher(traceConfig.rate), squeueFetcher: traceConfig.sharedFetcher, @@ -203,21 +191,22 @@ func (c *TraceCollector) uploadTrace(w http.ResponseWriter, r *http.Request) { } } -// detectTracePath returns the trace_root path based on the following criteria: +// detectTraceTemplatePath returns the trace_root path based on the following criteria: // 1. If TRACE_ROOT_PATH is specified, search that directory. If we don't find a templates dir, let's panic and crash the program. // 2. If TRACE_ROOT_PATH isn't specified, we can search cwd and /usr/share/prometheus-slurm-exporter. -func detectTracePath() string { - templateRootDir := "" - if path, ok := os.LookupEnv("TRACE_ROOT_PATH"); ok { - templateRootDir = path - if _, err := os.Stat(filepath.Join(templateRootDir, templateDirName)); err != nil { - panic("TRACE_ROOT_PATH must include a directory called: " + templateDirName) +// If no templates path is found, returns an empty string +func detectTraceTemplatePath() string { + if rpath, ok := os.LookupEnv("TRACE_ROOT_PATH"); ok { + templateP := filepath.Join(rpath, templateDirName) + if _, err := os.Stat(templateP); err != nil { + panic("TRACE_ROOT_PATH must include a directory called: templates") } - return templateRootDir + return templateP } - for _, p := range []string{".", "/usr/share/prometheus-slurm-exporter"} { - if _, err := os.Stat(filepath.Join(p, templateDirName)); err == nil { - return p + for _, rpath := range []string{".", "/usr/share/prometheus-slurm-exporter"} { + templateP := filepath.Join(rpath, templateDirName) + if _, err := os.Stat(templateP); err == nil { + return templateP } } return "" diff --git a/exporter/trace_test.go b/exporter/trace_test.go index 8244e03..d018e61 100644 --- a/exporter/trace_test.go +++ b/exporter/trace_test.go @@ -211,11 +211,11 @@ func TestDetectTraceRootPath_Env(t *testing.T) { testDir := t.TempDir() t.Setenv("TRACE_ROOT_PATH", testDir) // Ensure that the function panics if given a TRACE_ROOT_PATh with no 'templates' subdirectory - assert.PanicsWithValue(t, "TRACE_ROOT_PATH must include a directory called: templates", func() { detectTracePath() }) + assert.PanicsWithValue(t, "TRACE_ROOT_PATH must include a directory called: templates", func() { detectTraceTemplatePath() }) require.NoError(t, os.Mkdir(filepath.Join(testDir, templateDirName), 0o700)) // Now that we have a 'templates' subdir, it should no longer panic - assert.Equal(t, testDir, detectTracePath()) + assert.Equal(t, filepath.Join(testDir, templateDirName), detectTraceTemplatePath()) } func TestDetectTraceRootPath_Default(t *testing.T) { @@ -224,9 +224,9 @@ func TestDetectTraceRootPath_Default(t *testing.T) { os.Chdir(testDir) // Should come back empty if since we don't yet have a 'templates' subdir - assert.Equal(t, detectTracePath(), "") + assert.Equal(t, detectTraceTemplatePath(), "") require.NoError(t, os.Mkdir(filepath.Join(testDir, templateDirName), 0o700)) // Now that we have 'templates' subdir, cwd is a valid path - assert.Equal(t, detectTracePath(), ".") + assert.Equal(t, templateDirName, detectTraceTemplatePath()) }