Merge pull request #5 from poolski/refactor-and-add-tests

Fix bugs and add basic tests

.github/workflows/golangci-lint.yaml

@@ -18,10 +18,10 @@ jobs:
       - uses: actions/checkout@v3
       - uses: actions/setup-go@v4
         with:
-          go-version: "1.20"
+          go-version: "1.21"
           cache: false
       - name: golangci-lint
-        uses: golangci/golangci-lint-action@v3
+        uses: golangci/golangci-lint-action@v4
         with:
           version: latest
           # Optional: show only new issues if it's a pull request. The default value is `false`.

.github/workflows/release.yaml

@@ -7,46 +7,19 @@ permissions:
   packages: write
 
 jobs:
-  relese:
+  release:
     name: make releases
     runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        goos: [linux, darwin]
+        goarch: [amd64, arm64, arm]
+        goarm: [5, 6, 7]
     steps:
       - uses: actions/checkout@v3
       - uses: wangyoucao577/go-release-action@v1
         with:
           github_token: ${{ secrets.GITHUB_TOKEN }}
-          goos: linux
-          goarch: amd64
-      - uses: wangyoucao577/go-release-action@v1
-        with:
-          github_token: ${{ secrets.GITHUB_TOKEN }}
-          goos: linux
-          goarch: arm64
-      - uses: wangyoucao577/go-release-action@v1
-        with:
-          github_token: ${{ secrets.GITHUB_TOKEN }}
-          goos: linux
-          goarch: arm
-          goarm: 5
-      - uses: wangyoucao577/go-release-action@v1
-        with:
-          github_token: ${{ secrets.GITHUB_TOKEN }}
-          goos: linux
-          goarch: arm
-          goarm: 6
-      - uses: wangyoucao577/go-release-action@v1
-        with:
-          github_token: ${{ secrets.GITHUB_TOKEN }}
-          goos: linux
-          goarch: arm
-          goarm: 7
-      - uses: wangyoucao577/go-release-action@v1
-        with:
-          github_token: ${{ secrets.GITHUB_TOKEN }}
-          goos: darwin
-          goarch: amd64
-      - uses: wangyoucao577/go-release-action@v1
-        with:
-          github_token: ${{ secrets.GITHUB_TOKEN }}
-          goos: darwin
-          goarch: arm64
+          goos: ${{ matrix.goos }}
+          goarch: ${{ matrix.goarch }}
+          goarm: ${{ matrix.goarm }}

cmd/cmd.go → cmd/client/client.go

@@ -1,4 +1,4 @@
-package cmd
+package client
 
 import (
 	"crypto/tls"
@@ -14,43 +14,44 @@ import (
 	"github.com/spf13/viper"
 )
 
+type Config struct {
+	Days     int
+	Output   string
+	FilePath string
+	Insecure bool
+}
+
 type Client struct {
-	Conn      *websocket.Conn
-	MessageID int // MessageID is the ID of the message sent to the websocket. These must be incremented with each subsequent request
+	Config Config
+	Conn   *websocket.Conn
+	// MessageID represents the sequential ID of each message after the initial auth.
+	// These must be incremented with each subsequent request, otherwise the API will
+	// return an error.
+	MessageID int
 }
 
 // APIResponse represents the structure of the response received from the Home Assistant API.
 type APIResponse struct {
 	ID      int    `json:"id"`      // ID is the unique identifier of the response.
 	Type    string `json:"type"`    // Type is the type of the response.
 	Success bool   `json:"success"` // Success indicates whether the response was successful or not.
-	Result  struct {
-		ImportedElectricity []struct {
-			Start  int     `json:"start"`
-			End    int     `json:"end"`
-			Change float64 `json:"change"` // Change is the amount of electricity imported.
-		} `json:"sensor.smart_meter_electricity_import_2"`
+	Result  map[string][]struct {
+		Change float64 `json:"change"`
+		End    int64   `json:"end"`
+		Start  int64   `json:"start"`
 	} `json:"result"` // Result contains the data returned by the API.
 	Error struct {
 		Code    string `json:"code"`
 		Message string `json:"message"`
 	} `json:"error,omitempty"`
 }
 
-var (
-	days     int
-	output   string
-	csvFile  string
-	insecure bool
-)
-
 const hoursInADay = 24
 
-func init() {
-	rootCmd.PersistentFlags().IntVarP(&days, "days", "d", 30, "number of days to compute power stats for")
-	rootCmd.PersistentFlags().StringVarP(&output, "output", "o", "", "output format (text, table, csv)")
-	rootCmd.PersistentFlags().StringVarP(&csvFile, "csv-file", "f", "results.csv", "the path of the CSV file to write to")
-	rootCmd.PersistentFlags().BoolVarP(&insecure, "insecure", "i", false, "skip TLS verification")
+func New(cfg Config) *Client {
+	return &Client{
+		Config: cfg,
+	}
 }
 
 func (c *Client) Connect() error {
@@ -77,7 +78,7 @@ func (c *Client) Connect() error {
 	dialURL.Path = "/api/websocket"
 
 	// Skip TLS verification if insecure flag is set
-	if insecure {
+	if c.Config.Insecure {
 		dialer.TLSClientConfig = &tls.Config{
 			InsecureSkipVerify: true,
 		}
@@ -122,7 +123,7 @@ func (c *Client) Connect() error {
 // computePowerStats computes the power statistics for a given number of days and hours.
 // It prints a table to stdout where the rows are "days" and the columns are "hours".
 // The function writes the results to a CSV file and prints the averages to the console.
-func (c *Client) computePowerStats() {
+func (c *Client) ComputePowerStats() {
 	results, err := getResults(c)
 	if err != nil {
 		log.Error().Msg(fmt.Sprintf("getting results: %v", err))
@@ -136,7 +137,7 @@ func (c *Client) computePowerStats() {
 		for j := range results {
 			sum += results[j][i]
 		}
-		averages[i] = sum / float64(days)
+		averages[i] = sum / float64(c.Config.Days)
 	}
 
 	// Generate column headers for table/CSV
@@ -145,13 +146,13 @@ func (c *Client) computePowerStats() {
 		headers[i] = fmt.Sprintf("%d", i)
 	}
 
-	switch output {
+	switch c.Config.Output {
 	case "text":
 		writePlainText(averages)
 	case "table":
 		printTable(results, averages, headers)
 	case "csv":
-		err = writeCSVFile(headers, results, averages)
+		err = c.writeCSVFile(headers, results, averages)
 		if err != nil {
 			log.Error().Msg(fmt.Sprintf("writing CSV file: %v", err))
 			return
@@ -170,8 +171,8 @@ func writePlainText(averages []float64) {
 	}
 }
 
-func writeCSVFile(headers []string, results [][]float64, averages []float64) error {
-	f, err := os.Create(csvFile)
+func (c *Client) writeCSVFile(headers []string, results [][]float64, averages []float64) error {
+	f, err := os.Create(c.Config.FilePath)
 	if err != nil {
 		return fmt.Errorf("creating file: %w", err)
 	}
@@ -235,19 +236,18 @@ func getResults(c *Client) ([][]float64, error) {
 
 	// What we're doing is creating an offset from the current *day* based on a multiple of
 	// 24 hours, each time we iterate through the a "row" of the results slice.
-	results := make([][]float64, days)
+	results := make([][]float64, c.Config.Days)
+	sensorID := viper.GetString("sensor_id")
+	if sensorID == "" {
+		return nil, fmt.Errorf("sensor_id is required")
+	}
 
 	for i := range results {
 		c.MessageID++
 
 		offset := time.Duration((i+1)*24) * time.Hour
 		start := time.Now().Add(-offset).Truncate(24 * time.Hour).Format("2006-01-02T15:04:05.000Z")
 
-		sensorID := viper.GetString("sensor_id")
-		if sensorID == "" {
-			return nil, fmt.Errorf("sensor_id is required")
-		}
-
 		msg := map[string]interface{}{
 			"id":            c.MessageID,
 			"type":          "recorder/statistics_during_period",
@@ -274,15 +274,9 @@ func getResults(c *Client) ([][]float64, error) {
 		if !data.Success {
 			return nil, fmt.Errorf("api response error: %v", data.Error)
 		}
-
-		if len(data.Result.ImportedElectricity) != hoursInADay {
-			return nil, fmt.Errorf("expected %d sets of results, got %d", hoursInADay, len(data.Result.ImportedElectricity))
-		}
-
 		changeSlice := make([]float64, hoursInADay)
-		log.Debug().Msgf("got %d results", len(data.Result.ImportedElectricity))
 		for j := range changeSlice {
-			changeSlice[j] = data.Result.ImportedElectricity[j].Change
+			changeSlice[j] = data.Result[sensorID][j].Change
 		}
 		results[i] = changeSlice
 	}

cmd/client/client_test.go

@@ -0,0 +1,110 @@
+package client
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/gorilla/websocket"
+	"github.com/spf13/viper"
+	"gotest.tools/v3/assert"
+)
+
+func TestClient_Connect_SuccessfulConnection(t *testing.T) {
+	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// Handle the websocket upgrade request
+		upgrader := websocket.Upgrader{}
+		conn, _ := upgrader.Upgrade(w, r, nil)
+
+		// Read the initial message
+		initMsg := map[string]interface{}{
+			"type": "init",
+		}
+		assert.NilError(t, conn.WriteJSON(initMsg), "write initial message failed")
+
+		// Read the authentication message
+		var authMsg map[string]interface{}
+		assert.NilError(t, conn.ReadJSON(&authMsg), "read auth message failed")
+
+		// Check the authentication message
+		assert.Equal(t, authMsg["type"], "auth", "unexpected auth message type")
+		assert.Equal(t, authMsg["access_token"], "test_token", "unexpected access token")
+
+		// Send the authentication response
+		authResp := map[string]interface{}{
+			"type": "auth_ok",
+		}
+		assert.NilError(t, conn.WriteJSON(authResp), "write auth response failed")
+	}))
+	// Set up the client
+	client := &Client{
+		Config: Config{
+			Insecure: true,
+		},
+	}
+
+	// Set up the test environment
+	viper.Set("url", s.URL)
+	viper.Set("api_key", "test_token")
+
+	// Call the Connect method
+	err := client.Connect()
+
+	// Check the error condition
+	assert.NilError(t, err)
+}
+
+func TestClient_Connect_ErrorStates(t *testing.T) {
+	tests := []struct {
+		name     string
+		url      string
+		apiKey   string
+		expected string
+	}{
+		{
+			name:     "Empty URL",
+			url:      "",
+			apiKey:   "test_token",
+			expected: "url is required",
+		},
+		{
+			name:     "Invalid URL",
+			url:      "http://192.168.0.%31/",
+			apiKey:   "test_token",
+			expected: "parse \"http://192.168.0.%31/\": invalid URL escape \"%31\"",
+		},
+		{
+			name:     "Malformed URL",
+			url:      "htp:\\example.com",
+			apiKey:   "test_token",
+			expected: "dial: malformed ws or wss URL",
+		},
+		{
+			name:     "Bad Handshake",
+			url:      "http://example.com",
+			apiKey:   "test_token",
+			expected: "dial: websocket: bad handshake",
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			// Set up the client
+			client := &Client{
+				Config: Config{
+					Insecure: true,
+				},
+			}
+
+			// Set up the test environment
+			viper.Set("url", test.url)
+			viper.Set("api_key", test.apiKey)
+
+			// Call the Connect method
+			err := client.Connect()
+
+			// Check the error condition
+			assert.ErrorContains(t, err, test.expected)
+		})
+	}
+}

cmd/root.go

@@ -7,12 +7,20 @@ import (
 	"path/filepath"
 
 	"github.com/Songmu/prompter"
+	"github.com/poolski/powertracker/cmd/client"
 	"github.com/rs/zerolog/log"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
 )
 
-var cfgFile string
+var (
+	cfgFile string
+
+	days     int
+	output   string
+	csvFile  string
+	insecure bool
+)
 
 var rootCmd = &cobra.Command{
 	Use:   "powertracker",
@@ -24,11 +32,16 @@ var rootCmd = &cobra.Command{
 	It also saves the data to a CSV file in the current directory.`,
 
 	Run: func(cmd *cobra.Command, args []string) {
-		c := Client{}
+		c := client.New(client.Config{
+			Days:     days,
+			Output:   output,
+			FilePath: csvFile,
+			Insecure: insecure,
+		})
 		if err := c.Connect(); err != nil {
 			log.Fatal().Msgf("connecting to websocket: %s", err.Error())
 		}
-		c.computePowerStats()
+		c.ComputePowerStats()
 	},
 }
 
@@ -52,6 +65,11 @@ func init() {
 		sep := string(filepath.Separator)
 		confDir := home + sep + ".config"
 		rootCmd.PersistentFlags().StringVarP(&cfgFile, "config", "c", confDir+"/powertracker/config.yaml", "config file")
+
+		rootCmd.PersistentFlags().IntVarP(&days, "days", "d", 30, "number of days to compute power stats for")
+		rootCmd.PersistentFlags().StringVarP(&output, "output", "o", "", "output format (text, table, csv)")
+		rootCmd.PersistentFlags().StringVarP(&csvFile, "csv-file", "f", "results.csv", "the path of the CSV file to write to")
+		rootCmd.PersistentFlags().BoolVarP(&insecure, "insecure", "i", false, "skip TLS verification")
 	}
 }