ref: move server function to internal/

for easier integration testing
linode-obs · Oct 9, 2023 · a896e1a · a896e1a
1 parent e44d31c
commit a896e1a
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 27 deletions.
diff --git a/cmd/main.go b/cmd/main.go
@@ -6,23 +6,15 @@ import (
 	"net/http"
 	"os"
 
-	"github.com/prometheus/client_golang/prometheus/promhttp"
 	log "github.com/sirupsen/logrus"
-	"github.com/wbollock/ping_exporter/internal/collector"
+	"github.com/wbollock/ping_exporter/internal/server"
 )
 
 const (
-	version     = "0.1.0"
-	defaultHTML = `<html>
-			<head><title>Ping Exporter</title></head>
-			<body>
-			<h1>Ping Exporter</h1>
-			<p><a href='%s'>Metrics</a></p>
-			</body>
-			</html>`
+	version              = "0.1.0"
+	defaultLogLevel      = "info"
 	defaultListenAddress = ":9141"
 	defaultMetricsPath   = "/metrics"
-	defaultLogLevel      = "info"
 )
 
 var (
@@ -37,14 +29,6 @@ func printVersion() {
 	fmt.Printf("ping_exporter\nVersion: %s\nmulti-target ICMP prometheus exporter\n", version)
 }
 
-func mainPage(w http.ResponseWriter, r *http.Request) {
-	response := fmt.Sprintf(defaultHTML, *metricsPath)
-	_, err := w.Write([]byte(response))
-	if err != nil {
-		log.WithError(err).Error("Failed to write main page response")
-	}
-}
-
 func main() {
 	flag.Parse()
 
@@ -61,13 +45,10 @@ func main() {
 		log.SetLevel(log.InfoLevel)
 	}
 
-	log.Info("Listening on address ", *listenAddress)
-
-	http.Handle(*metricsPath, promhttp.Handler())
-	http.HandleFunc("/probe", collector.PingHandler)
-	http.HandleFunc("/", mainPage)
+	handler := server.SetupServer(*metricsPath)
 
-	if err := http.ListenAndServe(*listenAddress, nil); err != nil {
+	log.Infof("Starting server on %s", *listenAddress)
+	if err := http.ListenAndServe(*listenAddress, handler); err != nil {
 		log.WithError(err).Fatal("Failed to start the server")
 	}
 }
diff --git a/internal/collector/icmp_collector.go b/internal/collector/icmp_collector.go
@@ -147,7 +147,7 @@ func PingHandler(w http.ResponseWriter, r *http.Request) {
 	log.Debug("Request received with parameters ", p)
 
 	// TODO: ensure ResolveIPAddr is the best way to do lookups
-	// TODO: fails with IP address
+	// TODO: fails with IP address as target
 	ra, err := net.ResolveIPAddr(p.protocol, p.target)
 	if err != nil {
 		log.Error(err)
@@ -170,7 +170,7 @@ func PingHandler(w http.ResponseWriter, r *http.Request) {
 
 	if p.protocol == "icmp" {
 		pinger.SetPrivileged(true)
-	} else { // udp
+	} else {
 		pinger.SetPrivileged(false)
 	}
 

diff --git a/internal/server/server.go b/internal/server/server.go
@@ -0,0 +1,37 @@
+package server
+
+import (
+	"fmt"
+	"net/http"
+
+	"github.com/prometheus/client_golang/prometheus/promhttp"
+	log "github.com/sirupsen/logrus"
+	"github.com/wbollock/ping_exporter/internal/collector"
+)
+
+const (
+	defaultHTML = `<html>
+			<head><title>Ping Exporter</title></head>
+			<body>
+			<h1>Ping Exporter</h1>
+			<p><a href='%s'>Metrics</a></p>
+			</body>
+			</html>`
+	defaultMetricsPath = "/metrics"
+)
+
+func SetupServer(metricsPath string) http.Handler {
+	mux := http.NewServeMux()
+
+	mux.Handle(metricsPath, promhttp.Handler())
+	mux.HandleFunc("/probe", collector.PingHandler)
+	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+		response := fmt.Sprintf(defaultHTML, metricsPath)
+		_, err := w.Write([]byte(response))
+		if err != nil {
+			log.WithError(err).Error("Failed to write main page response")
+		}
+	})
+
+	return mux
+}