diff --git a/main.go b/main.go
index e3823d0..fb15f8b 100644
--- a/main.go
+++ b/main.go
@@ -15,7 +15,6 @@ import (
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"k8s.io/client-go/kubernetes"
 	"k8s.io/client-go/rest"
-	"k8s.io/client-go/tools/clientcmd"
 )
 
 var (
@@ -90,6 +89,23 @@ func (a *App) HandleSave(w http.ResponseWriter, r *http.Request) {
 var podIp = os.Getenv("POD_IP")
 var name = os.Getenv("POD_NAME")
 
+func GetDiscovery() Discovery {
+	if podIp == "" {
+		return nil
+	}
+
+	config, kerr := rest.InClusterConfig()
+
+	if kerr != nil {
+		log.Fatalf("Error creating kubernetes client: %v\n", kerr)
+	}
+	client, err := kubernetes.NewForConfig(config)
+	if err != nil {
+		log.Fatalf("Error creating client: %v\n", err)
+	}
+	return NewK8sDiscovery(client)
+}
+
 func main() {
 	// Create a new instance of the server
 	storage, err := NewDiskStorage(fmt.Sprintf("data/%s_state.gob", name))
@@ -101,30 +117,11 @@ func main() {
 		storage: storage,
 	}
 
-	var config *rest.Config
-	var kerr error
-	if podIp == "" {
-		config, kerr = clientcmd.BuildConfigFromFlags("", "/Users/mats/.kube/config")
-	} else {
-		config, kerr = rest.InClusterConfig()
-	}
-	if kerr != nil {
-		log.Fatalf("Error creating kubernetes client: %v\n", err)
-	}
-	client, err := kubernetes.NewForConfig(config)
-	if err != nil {
-		log.Fatalf("Error creating client: %v\n", err)
-	}
-	d := NewK8sDiscovery(client)
-
-	syncedPool, err := NewSyncedPool(app.pool, podIp, d)
+	syncedPool, err := NewSyncedPool(app.pool, podIp, GetDiscovery())
 	if err != nil {
 		log.Fatalf("Error creating synced pool: %v\n", err)
 	}
 
-	// if local
-	//syncedPool.AddRemote("localhost")
-
 	hg, err := NewGrainHandler(app.pool, ":1337")
 	if err != nil {
 		log.Fatalf("Error creating handler: %v\n", err)
diff --git a/message.go b/message.go
index b8a8d00..748a205 100644
--- a/message.go
+++ b/message.go
@@ -66,12 +66,9 @@ func (m Message) Write(w io.Writer) error {
 }
 
 func ReadMessage(reader io.Reader, m *Message) error {
-	h, err := GetMessageHandler(m.Type)
-	if err != nil {
-		return err
-	}
+
 	header := StorableMessageHeader{}
-	err = binary.Read(reader, binary.LittleEndian, &header)
+	err := binary.Read(reader, binary.LittleEndian, &header)
 	if err != nil {
 		return err
 	}
@@ -80,7 +77,10 @@ func ReadMessage(reader io.Reader, m *Message) error {
 	if err != nil {
 		return err
 	}
-
+	h, err := GetMessageHandler(header.Type)
+	if err != nil {
+		return err
+	}
 	content, err := h.Read(messageBytes)
 	if err != nil {
 		return err