Refactor MQTT connect / disconnect

1 files changed, 29 insertions, 9 deletions

diff --git a/main.go b/main.go
index 075feb8..928ae2c 100644
--- a/main.go
+++ b/main.go
@@ -3,6 +3,7 @@ package main
 import (
 	"fmt"
 	"log"
+	"time"
 
 	"xengineering.eu/homematic-go/homematic"
 	mqtt "github.com/eclipse/paho.mqtt.golang"
@@ -15,6 +16,8 @@ const (
 	TOPIC_PREFIX = `sia`
 	QOS = byte(0)
 	RETAINED = false
+	MQTT_CONNECT_TIMEOUT = time.Second * 5
+	MQTT_DISCONNECT_TIMEOUT_US = 500
 )
 
 func main() {
@@ -25,15 +28,11 @@ func main() {
 		log.Fatalf("Failed to retrieve device list: %v", err)
 	}
 
-	opts := mqtt.NewClientOptions()
-	opts.AddBroker(BROKER)
-	opts.SetClientID(CLIENT_ID)
-	opts.SetCleanSession(true)
-
-	client := mqtt.NewClient(opts)
-	if token := client.Connect(); token.Wait() && token.Error() != nil {
-		log.Fatal(token.Error())
+	client, err := ConnectMQTT(BROKER, CLIENT_ID)
+	if err != nil {
+		log.Fatalf("Could not connect to MQTT broker: %v", err)
 	}
+	defer client.Disconnect(MQTT_DISCONNECT_TIMEOUT_US)
 
 	for _, device := range inventory {
 		if device.Type == `SHUTTER_CONTACT` {
@@ -47,6 +46,27 @@ func main() {
 			client.Publish(topic, QOS, RETAINED, payload)
 		}
 	}
+}
+
+func ConnectMQTT(broker string, id string) (mqtt.Client, error) {
+	opts := mqtt.NewClientOptions()
+	opts.AddBroker(broker)
+	opts.SetClientID(id)
+	opts.SetCleanSession(true)
+
+	client := mqtt.NewClient(opts)
+
+	token := client.Connect()
+
+	success := token.WaitTimeout(MQTT_CONNECT_TIMEOUT)
+	if !success {
+		return client, fmt.Errorf("Timed out after %v.", MQTT_CONNECT_TIMEOUT)
+	}
+
+	err := token.Error()
+	if err != nil {
+		return client, err
+	}
 
-	client.Disconnect(500)
+	return client, nil
 }