Move prober and config to own package

config.go → config/config.go

@@ -1,12 +1,15 @@
-package main
+package config
 
 import (
 	"errors"
 	"fmt"
+	"io/ioutil"
 	"strings"
 	"sync"
 	"time"
 
+	yaml "gopkg.in/yaml.v2"
+
 	"github.com/prometheus/common/config"
 )
 
@@ -22,6 +25,25 @@ type SafeConfig struct {
 	C *Config
 }
 
+func (sc *SafeConfig) ReloadConfig(confFile string) (err error) {
+	var c = &Config{}
+
+	yamlFile, err := ioutil.ReadFile(confFile)
+	if err != nil {
+		return fmt.Errorf("Error reading config file: %s", err)
+	}
+
+	if err := yaml.Unmarshal(yamlFile, c); err != nil {
+		return fmt.Errorf("Error parsing config file: %s", err)
+	}
+
+	sc.Lock()
+	sc.C = c
+	sc.Unlock()
+
+	return nil
+}
+
 type Module struct {
 	Prober  string        `yaml:"prober,omitempty"`
 	Timeout time.Duration `yaml:"timeout,omitempty"`

config_test.go → config/config_test.go

@@ -1,4 +1,4 @@
-package main
+package config
 
 import (
 	"strings"
@@ -12,7 +12,7 @@ func TestLoadConfig(t *testing.T) {
 		C: &Config{},
 	}
 
-	err := sc.reloadConfig("testdata/blackbox-good.yml")
+	err := sc.ReloadConfig("testdata/blackbox-good.yml")
 	if err != nil {
 		t.Errorf("Error loading config %v: %v", "blackbox.yml", err)
 	}
@@ -28,15 +28,15 @@ func TestLoadBadConfigs(t *testing.T) {
 	}{
 		{
 			ConfigFile:    "testdata/blackbox-bad.yml",
-			ExpectedError: "unknown fields in dns probe: invalid_extra_field",
+			ExpectedError: "Error parsing config file: unknown fields in dns probe: invalid_extra_field",
 		},
 		{
 			ConfigFile:    "testdata/invalid-dns-module.yml",
-			ExpectedError: "Query name must be set for DNS module",
+			ExpectedError: "Error parsing config file: Query name must be set for DNS module",
 		},
 	}
 	for i, test := range tests {
-		err := sc.reloadConfig(test.ConfigFile)
+		err := sc.ReloadConfig(test.ConfigFile)
 		if err.Error() != test.ExpectedError {
 			t.Errorf("In case %v:\nExpected:\n%v\nGot:\n%v", i, test.ExpectedError, err.Error())
 		}
@@ -48,7 +48,7 @@ func TestHideConfigSecrets(t *testing.T) {
 		C: &Config{},
 	}
 
-	err := sc.reloadConfig("testdata/blackbox-good.yml")
+	err := sc.ReloadConfig("testdata/blackbox-good.yml")
 	if err != nil {
 		t.Errorf("Error loading config %v: %v", "testdata/blackbox-good.yml", err)
 	}

main.go

@@ -16,7 +16,6 @@ package main
 import (
 	"context"
 	"fmt"
-	"io/ioutil"
 	"net/http"
 	"os"
 	"os/signal"
@@ -31,48 +30,29 @@ import (
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"github.com/prometheus/common/log"
 	"github.com/prometheus/common/version"
+
+	"github.com/prometheus/blackbox_exporter/config"
+	"github.com/prometheus/blackbox_exporter/prober"
 )
 
 var (
-	sc = &SafeConfig{
-		C: &Config{},
+	sc = &config.SafeConfig{
+		C: &config.Config{},
 	}
 
 	configFile    = kingpin.Flag("config.file", "Blackbox exporter configuration file.").Default("blackbox.yml").String()
 	listenAddress = kingpin.Flag("web.listen-address", "The address to listen on for HTTP requests.").Default(":9115").String()
 	timeoutOffset = kingpin.Flag("timeout-offset", "Offset to subtract from timeout in seconds.").Default("0.5").Float64()
-)
-
-var Probers = map[string]func(context.Context, string, Module, *prometheus.Registry) bool{
-	"http": probeHTTP,
-	"tcp":  probeTCP,
-	"icmp": probeICMP,
-	"dns":  probeDNS,
-}
-
-func (sc *SafeConfig) reloadConfig(confFile string) (err error) {
-	var c = &Config{}
-
-	yamlFile, err := ioutil.ReadFile(confFile)
-	if err != nil {
-		log.Errorf("Error reading config file: %s", err)
-		return err
-	}
 
-	if err := yaml.Unmarshal(yamlFile, c); err != nil {
-		log.Errorf("Error parsing config file: %s", err)
-		return err
+	Probers = map[string]prober.ProbeFn{
+		"http": prober.ProbeHTTP,
+		"tcp":  prober.ProbeTCP,
+		"icmp": prober.ProbeICMP,
+		"dns":  prober.ProbeDNS,
 	}
+)
 
-	sc.Lock()
-	sc.C = c
-	sc.Unlock()
-
-	log.Infoln("Loaded config file")
-	return nil
-}
-
-func probeHandler(w http.ResponseWriter, r *http.Request, c *Config) {
+func probeHandler(w http.ResponseWriter, r *http.Request, c *config.Config) {
 	moduleName := r.URL.Query().Get("module")
 	if moduleName == "" {
 		moduleName = "http_2xx"
@@ -151,9 +131,10 @@ func main() {
 	log.Infoln("Starting blackbox_exporter", version.Info())
 	log.Infoln("Build context", version.BuildContext())
 
-	if err := sc.reloadConfig(*configFile); err != nil {
+	if err := sc.ReloadConfig(*configFile); err != nil {
 		log.Fatalf("Error loading config: %s", err)
 	}
+	log.Infoln("Loaded config file")
 
 	hup := make(chan os.Signal)
 	reloadCh := make(chan chan error)
@@ -162,14 +143,17 @@ func main() {
 		for {
 			select {
 			case <-hup:
-				if err := sc.reloadConfig(*configFile); err != nil {
+				if err := sc.ReloadConfig(*configFile); err != nil {
 					log.Errorf("Error reloading config: %s", err)
+					continue
 				}
+				log.Infoln("Loaded config file")
 			case rc := <-reloadCh:
-				if err := sc.reloadConfig(*configFile); err != nil {
+				if err := sc.ReloadConfig(*configFile); err != nil {
 					log.Errorf("Error reloading config: %s", err)
 					rc <- err
 				} else {
+					log.Infoln("Loaded config file")
 					rc <- nil
 				}
 			}

main_test.go

@@ -0,0 +1,44 @@
+package main
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/prometheus/blackbox_exporter/config"
+)
+
+var c = &config.Config{
+	Modules: map[string]config.Module{
+		"http_2xx": config.Module{
+			Prober:  "http",
+			Timeout: 10 * time.Second,
+		},
+	},
+}
+
+func TestPrometheusTimeoutHTTP(t *testing.T) {
+
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		time.Sleep(2 * time.Second)
+	}))
+	defer ts.Close()
+
+	req, err := http.NewRequest("GET", "?target="+ts.URL, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	req.Header.Set("X-Prometheus-Scrape-Timeout-Seconds", "1")
+
+	rr := httptest.NewRecorder()
+	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		probeHandler(w, r, c)
+	})
+
+	handler.ServeHTTP(rr, req)
+
+	if status := rr.Code; status != http.StatusOK {
+		t.Errorf("probe request handler returned wrong status code: %v, want %v", status, http.StatusOK)
+	}
+}

dns.go → prober/dns.go

@@ -11,7 +11,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package main
+package prober
 
 import (
 	"context"
@@ -22,10 +22,12 @@ import (
 	"github.com/miekg/dns"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/common/log"
+
+	"github.com/prometheus/blackbox_exporter/config"
 )
 
 // validRRs checks a slice of RRs received from the server against a DNSRRValidator.
-func validRRs(rrs *[]dns.RR, v *DNSRRValidator) bool {
+func validRRs(rrs *[]dns.RR, v *config.DNSRRValidator) bool {
 	// Fail the probe if there are no RRs of a given type, but a regexp match is required
 	// (i.e. FailIfNotMatchesRegexp is set).
 	if len(*rrs) == 0 && len(v.FailIfNotMatchesRegexp) > 0 {
@@ -82,7 +84,7 @@ func validRcode(rcode int, valid []string) bool {
 	return false
 }
 
-func probeDNS(ctx context.Context, target string, module Module, registry *prometheus.Registry) bool {
+func ProbeDNS(ctx context.Context, target string, module config.Module, registry *prometheus.Registry) bool {
 	var numAnswer, numAuthority, numAdditional int
 	var dialProtocol string
 	probeDNSAnswerRRSGauge := prometheus.NewGauge(prometheus.GaugeOpts{