refactor(battery): move logic to environment

src/environment/shell.go

@@ -9,6 +9,7 @@ import (
 	"io/fs"
 	"io/ioutil"
 	"log"
+	"math"
 	"net/http"
 	"oh-my-posh/regex"
 	"os"
@@ -122,6 +123,11 @@ type TemplateCache struct {
 	WSL          bool
 }
 
+type BatteryInfo struct {
+	Percentage int
+	State      battery.State
+}
+
 type Environment interface {
 	Getenv(key string) string
 	Pwd() string
@@ -149,7 +155,7 @@ type Environment interface {
 	RunShellCommand(shell, command string) string
 	ExecutionTime() float64
 	Flags() *Flags
-	BatteryInfo() ([]*battery.Battery, error)
+	BatteryState() (*BatteryInfo, error)
 	QueryWindowTitles(processName, windowTitleRegex string) (string, error)
 	WindowsRegistryKeyValue(path string) (*WindowsRegistryValue, error)
 	HTTPRequest(url string, timeout int, requestModifiers ...HTTPRequestModifier) ([]byte, error)
@@ -560,8 +566,43 @@ func (env *ShellEnvironment) Flags() *Flags {
 	return env.CmdFlags
 }
 
-func (env *ShellEnvironment) BatteryInfo() ([]*battery.Battery, error) {
+func mapMostLogicalState(currentState, newState battery.State) battery.State {
+	switch currentState {
+	case battery.Discharging, battery.NotCharging:
+		return battery.Discharging
+	case battery.Empty:
+		return newState
+	case battery.Charging:
+		if newState == battery.Discharging {
+			return battery.Discharging
+		}
+		return battery.Charging
+	case battery.Unknown:
+		return newState
+	case battery.Full:
+		return newState
+	}
+	return newState
+}
+
+func (env *ShellEnvironment) BatteryState() (*BatteryInfo, error) {
 	defer env.trace(time.Now(), "BatteryInfo")
+
+	parseBatteryInfo := func(batteries []*battery.Battery) *BatteryInfo {
+		var info BatteryInfo
+		var current, total float64
+		var state battery.State
+		for _, bt := range batteries {
+			current += bt.Current
+			total += bt.Full
+			state = mapMostLogicalState(state, bt.State)
+		}
+		batteryPercentage := current / total * 100
+		info.Percentage = int(math.Min(100, batteryPercentage))
+		info.State = state
+		return &info
+	}
+
 	batteries, err := battery.GetAll()
 	// actual error, return it
 	if err != nil && len(batteries) == 0 {
@@ -609,7 +650,7 @@ func (env *ShellEnvironment) BatteryInfo() ([]*battery.Battery, error) {
 
 	// when battery info fails to get retrieved but there is at least one valid battery, return it without error
 	if len(validBatteries) > 0 && fatalErr != nil && strings.Contains(fatalErr.Error(), unableToRetrieveBatteryInfo) {
-		return validBatteries, nil
+		return parseBatteryInfo(validBatteries), nil
 	}
 	// another error occurred (possibly unmapped use-case), return it
 	if fatalErr != nil {
@@ -617,7 +658,7 @@ func (env *ShellEnvironment) BatteryInfo() ([]*battery.Battery, error) {
 		return nil, fatalErr
 	}
 	// everything is fine
-	return validBatteries, nil
+	return parseBatteryInfo(validBatteries), nil
 }
 
 func (env *ShellEnvironment) Shell() string {

src/environment/shell_test.go

@@ -3,6 +3,7 @@ package environment
 import (
 	"testing"
 
+	"github.com/distatus/battery"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -97,3 +98,24 @@ func TestDirMatchesOneOfRegexInvertedNonEscaped(t *testing.T) {
 	}()
 	_ = dirMatchesOneOf("Projects/oh-my-posh", "", LinuxPlatform, []string{"(?!Projects/).*"})
 }
+
+func TestMapBatteriesState(t *testing.T) {
+	cases := []struct {
+		Case          string
+		ExpectedState battery.State
+		CurrentState  battery.State
+		NewState      battery.State
+	}{
+		{Case: "charging > charged", ExpectedState: battery.Charging, CurrentState: battery.Full, NewState: battery.Charging},
+		{Case: "charging < discharging", ExpectedState: battery.Discharging, CurrentState: battery.Discharging, NewState: battery.Charging},
+		{Case: "charging == charging", ExpectedState: battery.Charging, CurrentState: battery.Charging, NewState: battery.Charging},
+		{Case: "discharging > charged", ExpectedState: battery.Discharging, CurrentState: battery.Full, NewState: battery.Discharging},
+		{Case: "discharging > unknown", ExpectedState: battery.Discharging, CurrentState: battery.Unknown, NewState: battery.Discharging},
+		{Case: "discharging > full", ExpectedState: battery.Discharging, CurrentState: battery.Full, NewState: battery.Discharging},
+		{Case: "discharging > charging 2", ExpectedState: battery.Discharging, CurrentState: battery.Charging, NewState: battery.Discharging},
+		{Case: "discharging > empty", ExpectedState: battery.Discharging, CurrentState: battery.Empty, NewState: battery.Discharging},
+	}
+	for _, tc := range cases {
+		assert.Equal(t, tc.ExpectedState, mapMostLogicalState(tc.CurrentState, tc.NewState), tc.Case)
+	}
+}

src/mock/environment.go

@@ -4,7 +4,6 @@ import (
 	"io/fs"
 	"oh-my-posh/environment"
 
-	"github.com/distatus/battery"
 	mock "github.com/stretchr/testify/mock"
 )
 
@@ -122,9 +121,9 @@ func (env *MockedEnvironment) Flags() *environment.Flags {
 	return arguments.Get(0).(*environment.Flags)
 }
 
-func (env *MockedEnvironment) BatteryInfo() ([]*battery.Battery, error) {
+func (env *MockedEnvironment) BatteryState() (*environment.BatteryInfo, error) {
 	args := env.Called()
-	return args.Get(0).([]*battery.Battery), args.Error(1)
+	return args.Get(0).(*environment.BatteryInfo), args.Error(1)
 }
 
 func (env *MockedEnvironment) Shell() string {

src/segments/battery.go

@@ -1,7 +1,6 @@
 package segments
 
 import (
-	"math"
 	"oh-my-posh/environment"
 	"oh-my-posh/properties"
 
@@ -12,8 +11,7 @@ type Battery struct {
 	props properties.Properties
 	env   environment.Environment
 
-	battery.Battery
+	*environment.BatteryInfo
-	Percentage int
 	Error string
 	Icon  string
 }
@@ -37,26 +35,19 @@ func (b *Battery) Enabled() bool {
 		return false
 	}
 
-	batteries, err := b.env.BatteryInfo()
+	var err error
+	b.BatteryInfo, err = b.env.BatteryState()
 
 	if !b.enabledWhileError(err) {
 		return false
 	}
 
 	// case on computer without batteries(no error, empty array)
-	if err == nil && len(batteries) == 0 {
+	if err == nil && b.BatteryInfo == nil {
 		return false
 	}
 
-	for _, bt := range batteries {
+	switch b.BatteryInfo.State {
-		b.Battery.Current += bt.Current
-		b.Battery.Full += bt.Full
-		b.Battery.State = b.mapMostLogicalState(b.Battery.State, bt.State)
-	}
-	batteryPercentage := b.Battery.Current / b.Battery.Full * 100
-	b.Percentage = int(math.Min(100, batteryPercentage))
-
-	switch b.Battery.State {
 	case battery.Discharging, battery.NotCharging:
 		b.Icon = b.props.GetString(DischargingIcon, "")
 	case battery.Charging:
@@ -85,31 +76,11 @@ func (b *Battery) enabledWhileError(err error) bool {
 	// This hack ensures we display a fully charged battery, even if
 	// that state can be incorrect. It's better to "ignore" the error
 	// than to not display the segment at all as that will confuse users.
-	b.Battery.Current = 100
+	b.Percentage = 100
-	b.Battery.Full = 10
+	b.State = battery.Full
-	b.Battery.State = battery.Full
 	return true
 }
 
-func (b *Battery) mapMostLogicalState(currentState, newState battery.State) battery.State {
-	switch currentState {
-	case battery.Discharging, battery.NotCharging:
-		return battery.Discharging
-	case battery.Empty:
-		return newState
-	case battery.Charging:
-		if newState == battery.Discharging {
-			return battery.Discharging
-		}
-		return battery.Charging
-	case battery.Unknown:
-		return newState
-	case battery.Full:
-		return newState
-	}
-	return newState
-}
-
 func (b *Battery) Init(props properties.Properties, env environment.Environment) {
 	b.props = props
 	b.env = env

src/segments/battery_test.go

@@ -1,30 +0,0 @@
-package segments
-
-import (
-	"testing"
-
-	"github.com/distatus/battery"
-	"github.com/stretchr/testify/assert"
-)
-
-func TestMapBatteriesState(t *testing.T) {
-	cases := []struct {
-		Case          string
-		ExpectedState battery.State
-		CurrentState  battery.State
-		NewState      battery.State
-	}{
-		{Case: "charging > charged", ExpectedState: battery.Charging, CurrentState: battery.Full, NewState: battery.Charging},
-		{Case: "charging < discharging", ExpectedState: battery.Discharging, CurrentState: battery.Discharging, NewState: battery.Charging},
-		{Case: "charging == charging", ExpectedState: battery.Charging, CurrentState: battery.Charging, NewState: battery.Charging},
-		{Case: "discharging > charged", ExpectedState: battery.Discharging, CurrentState: battery.Full, NewState: battery.Discharging},
-		{Case: "discharging > unknown", ExpectedState: battery.Discharging, CurrentState: battery.Unknown, NewState: battery.Discharging},
-		{Case: "discharging > full", ExpectedState: battery.Discharging, CurrentState: battery.Full, NewState: battery.Discharging},
-		{Case: "discharging > charging 2", ExpectedState: battery.Discharging, CurrentState: battery.Charging, NewState: battery.Discharging},
-		{Case: "discharging > empty", ExpectedState: battery.Discharging, CurrentState: battery.Empty, NewState: battery.Discharging},
-	}
-	for _, tc := range cases {
-		batt := &Battery{}
-		assert.Equal(t, tc.ExpectedState, batt.mapMostLogicalState(tc.CurrentState, tc.NewState), tc.Case)
-	}
-}