add the NewPidFileFn to helper

Signed-off-by: sbookworm <lisong.cruise@gmail.com>
This commit is contained in:
sbookworm 2020-10-10 16:37:17 +08:00
parent 47cfdc9bb8
commit 469ec2747b
2 changed files with 85 additions and 0 deletions

View File

@ -15,7 +15,11 @@ package prometheus
import ( import (
"errors" "errors"
"fmt"
"io/ioutil"
"os" "os"
"strconv"
"strings"
) )
type processCollector struct { type processCollector struct {
@ -149,3 +153,20 @@ func (c *processCollector) reportError(ch chan<- Metric, desc *Desc, err error)
} }
ch <- NewInvalidMetric(desc, err) ch <- NewInvalidMetric(desc, err)
} }
// NewPidFileFn returns a function that retrieves a pid from the specified file.
// It is meant to be used for the PidFn field in ProcessCollectorOpts.
func NewPidFileFn(pidFilePath string) func() (int, error) {
return func() (int, error) {
content, err := ioutil.ReadFile(pidFilePath)
if err != nil {
return 0, fmt.Errorf("can't read pid file %q: %+v", pidFilePath, err)
}
pid, err := strconv.Atoi(strings.TrimSpace(string(content)))
if err != nil {
return 0, fmt.Errorf("can't parse pid file %q: %+v", pidFilePath, err)
}
return pid, nil
}
}

View File

@ -18,8 +18,11 @@ package prometheus
import ( import (
"bytes" "bytes"
"errors" "errors"
"fmt"
"os" "os"
"path/filepath"
"regexp" "regexp"
"strings"
"testing" "testing"
"github.com/prometheus/common/expfmt" "github.com/prometheus/common/expfmt"
@ -101,3 +104,64 @@ func TestProcessCollector(t *testing.T) {
t.Errorf("%d metrics collected, want 1", n) t.Errorf("%d metrics collected, want 1", n)
} }
} }
func TestNewPidFileFn(t *testing.T) {
folderPath, err := os.Getwd()
if err != nil {
t.Error("failed to get current path")
}
mockPidFilePath := filepath.Join(folderPath, "mockPidFile")
defer os.Remove(mockPidFilePath)
testCases := []struct {
mockPidFile func()
expectedErrPrefix string
expectedPid int
desc string
}{
{
mockPidFile: func() {
os.Remove(mockPidFilePath)
},
expectedErrPrefix: "can't read pid file",
expectedPid: 0,
desc: "no existed pid file",
},
{
mockPidFile: func() {
os.Remove(mockPidFilePath)
f, _ := os.Create(mockPidFilePath)
f.Write([]byte("abc"))
f.Close()
},
expectedErrPrefix: "can't parse pid file",
expectedPid: 0,
desc: "existed pid file, error pid number",
},
{
mockPidFile: func() {
os.Remove(mockPidFilePath)
f, _ := os.Create(mockPidFilePath)
f.Write([]byte("123"))
f.Close()
},
expectedErrPrefix: "",
expectedPid: 123,
desc: "existed pid file, correct pid number",
},
}
for _, tc := range testCases {
fn := NewPidFileFn(mockPidFilePath)
if fn == nil {
t.Error("Should not get nil PidFileFn")
}
tc.mockPidFile()
if pid, err := fn(); pid != tc.expectedPid || (err != nil && !strings.HasPrefix(err.Error(), tc.expectedErrPrefix)) {
fmt.Println(err.Error())
t.Error(tc.desc)
}
}
}