From b044de125d907b5a3abd37171c0cb3d02f4d6a48 Mon Sep 17 00:00:00 2001
From: Cody Lee <cody.lee@datadoghq.com>
Date: Fri, 21 Jul 2023 11:36:24 -0500
Subject: [PATCH] plug in mock status response and improved logging

---
 mocks/mock_client.go |  6 +++---
 mocks/mock_server.go |  9 +++++++--
 types.go             |  2 +-
 unifi.go             | 11 ++++++++---
 4 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/mocks/mock_client.go b/mocks/mock_client.go
index 0f9216b..a9128e2 100644
--- a/mocks/mock_client.go
+++ b/mocks/mock_client.go
@@ -366,16 +366,16 @@ func (m *MockUnifi) Logout() error {
 
 // GetServerData sets the controller's version and UUID. Only call this if you
 // previously called Login and suspect the controller version has changed.
-func (m *MockUnifi) GetServerData() error {
+func (m *MockUnifi) GetServerData() (*unifi.ServerStatus, error) {
 	var response struct {
 		Data unifi.ServerStatus `json:"meta"`
 	}
 
 	err := m.faker.Struct(&response)
 	if err != nil {
-		return err
+		return nil, err
 	}
-	return nil
+	return &response.Data, nil
 }
 
 // GetUsers returns a response full of clients that connected to the UDM within the provided amount of time
diff --git a/mocks/mock_server.go b/mocks/mock_server.go
index 21bf80e..d48fe87 100644
--- a/mocks/mock_server.go
+++ b/mocks/mock_server.go
@@ -3,6 +3,7 @@ package mocks
 import (
 	"encoding/json"
 	"fmt"
+	"log"
 	"net/http"
 	"net/http/httptest"
 	"regexp"
@@ -54,6 +55,7 @@ var (
 )
 
 func respondResultOrErr(w http.ResponseWriter, v any, err error) {
+	log.Printf("[DEBUG] Answering mock response err=%+v value=%+v\n", err, v)
 	if err != nil {
 		b, _ := json.Marshal(err)
 		w.WriteHeader(500)
@@ -67,14 +69,16 @@ func respondResultOrErr(w http.ResponseWriter, v any, err error) {
 
 func (m *MockHTTPTestServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	p := strings.TrimSpace(r.URL.Path)
+	log.Printf("[DEBUG] Received mock request path=%s\n", p)
 	switch {
 	case apiRogueAP.MatchString(p):
 		aps, err := m.mocked.GetRogueAPs(nil)
 		respondResultOrErr(w, aps, err)
 		return
 	case apiStatusPath.MatchString(p):
-		
-		// todo
+		s, err := m.mocked.GetServerData()
+		respondResultOrErr(w, s, err)
+		return
 	case apiEventPath.MatchString(p):
 		events, err := m.mocked.GetEvents(nil, time.Hour)
 		respondResultOrErr(w, events, err)
@@ -141,6 +145,7 @@ func (m *MockHTTPTestServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		w.WriteHeader(501)
 		return
 	default:
+		log.Println("[DEBUG] Answering mock response err=404 not found")
 		http.NotFoundHandler().ServeHTTP(w, r)
 		return
 	}
diff --git a/types.go b/types.go
index 9edd412..37c6fad 100644
--- a/types.go
+++ b/types.go
@@ -209,7 +209,7 @@ type UnifiClient interface {
 	Logout() error
 	// GetServerData sets the controller's version and UUID. Only call this if you
 	// previously called Login and suspect the controller version has changed.
-	GetServerData() error
+	GetServerData() (*ServerStatus, error)
 	// GetUsers returns a response full of clients that connected to the UDM within the provided amount of time
 	// using the insight historical connection data set.
 	GetUsers(sites []*Site, hours int) ([]*User, error)
diff --git a/unifi.go b/unifi.go
index 48aa313..0100fbf 100644
--- a/unifi.go
+++ b/unifi.go
@@ -55,7 +55,7 @@ func NewUnifi(config *Config) (*Unifi, error) {
 		return u, err
 	}
 
-	if err := u.GetServerData(); err != nil {
+	if _, err := u.GetServerData(); err != nil {
 		return u, fmt.Errorf("unable to get server version: %w", err)
 	}
 
@@ -201,14 +201,19 @@ func (u *Unifi) checkNewStyleAPI() error {
 
 // GetServerData sets the controller's version and UUID. Only call this if you
 // previously called Login and suspect the controller version has changed.
-func (u *Unifi) GetServerData() error {
+func (u *Unifi) GetServerData() (*ServerStatus, error) {
 	var response struct {
 		Data ServerStatus `json:"meta"`
 	}
 
+	err := u.GetData(APIStatusPath, &response)
+	if err != nil {
+		return nil, err
+	}
+
 	u.ServerStatus = &response.Data
 
-	return u.GetData(APIStatusPath, &response)
+	return u.ServerStatus, nil
 }
 
 // GetData makes a unifi request and unmarshals the response into a provided pointer.