Skip to content

Commit 8888946

Browse files
committed
fix: raise on init
Signed-off-by: Martin Buchleitner <mabunixda@gmail.com>
1 parent 40346a4 commit 8888946

File tree

3 files changed

+59
-15
lines changed

3 files changed

+59
-15
lines changed

helper.go

Lines changed: 36 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,11 @@ import (
55
"encoding/hex"
66
"fmt"
77
"math/rand"
8+
"runtime"
9+
810
"time"
911
)
1012

11-
var randomSource = rand.New(rand.NewSource(time.Now().UnixNano()))
12-
1313
func Keys[K comparable, V any](m map[K]V) []K {
1414
keys := make([]K, 0, len(m))
1515
for k := range m {
@@ -29,10 +29,43 @@ func sha256sum(data string) string {
2929
}
3030
func randomHexString(n int) string {
3131
b := make([]byte, (n+2)/2) // can be simplified to n/2 if n is always even
32-
32+
randomSource := rand.New(rand.NewSource(time.Now().UnixNano()))
3333
if _, err := randomSource.Read(b); err != nil {
3434
panic(err)
3535
}
3636

3737
return hex.EncodeToString(b)[1 : n+1]
3838
}
39+
40+
// func getCurrentFunctionName() string {
41+
// // Skip GetCurrentFunctionName
42+
// return getFrame(1).Function
43+
// }
44+
45+
func getCallerFunctionName() string {
46+
// Skip GetCallerFunctionName and the function to get the caller of
47+
return getFrame(2).Function
48+
}
49+
50+
func getFrame(skipFrames int) runtime.Frame {
51+
// We need the frame at index skipFrames+2, since we never want runtime.Callers and getFrame
52+
targetFrameIndex := skipFrames + 2
53+
54+
// Set size to targetFrameIndex+2 to ensure we have room for one more caller than we need
55+
programCounters := make([]uintptr, targetFrameIndex+2)
56+
n := runtime.Callers(0, programCounters)
57+
58+
frame := runtime.Frame{Function: "unknown"}
59+
if n > 0 {
60+
frames := runtime.CallersFrames(programCounters[:n])
61+
for more, frameIndex := true, 0; more && frameIndex <= targetFrameIndex; frameIndex++ {
62+
var frameCandidate runtime.Frame
63+
frameCandidate, more = frames.Next()
64+
if frameIndex == targetFrameIndex {
65+
frame = frameCandidate
66+
}
67+
}
68+
}
69+
70+
return frame
71+
}

wattpilot.go

Lines changed: 21 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -25,8 +25,9 @@ import (
2525
)
2626

2727
const (
28-
ContextTimeout = 60 // seconds
29-
ReconnectTimeout = 5 // seconds
28+
ContextTimeout = 60 // seconds
29+
ConnectionTimeout = 30 // seconds
30+
ReconnectTimeout = 5 // seconds
3031

3132
EventTypeHello = "hello"
3233
EventTypeAuthRequired = "authRequired"
@@ -401,7 +402,7 @@ func (w *Wattpilot) Connect() error {
401402
go w.connectionManager()
402403
})
403404

404-
if w.connectAndWait(30 * time.Second) {
405+
if w.connectAndWait(ConnectionTimeout * time.Second) {
405406
return nil
406407
}
407408

@@ -429,7 +430,7 @@ func (w *Wattpilot) connectImpl() error {
429430
return nil
430431
}
431432

432-
ctx, cancel := context.WithTimeout(context.Background(), ContextTimeout*time.Second)
433+
ctx, cancel := context.WithTimeout(context.Background(), ConnectionTimeout*time.Second)
433434
defer cancel()
434435

435436
conn, _, err := websocket.Dial(ctx, fmt.Sprintf("ws://%s/ws", w.host), nil)
@@ -444,21 +445,28 @@ func (w *Wattpilot) connectImpl() error {
444445
select {
445446
case w.isConnected = <-w.connected:
446447
if !w.isConnected {
448+
w.logger.WithFields(logrus.Fields{"wattpilot": w.host}).Info("Authentication timeout")
447449
w.disconnectImpl()
448450
return errors.New("authentication failed")
449451
}
450452
case <-ctx.Done():
453+
w.logger.WithFields(logrus.Fields{"wattpilot": w.host}).Info("Connection handshake timeout")
451454
w.disconnectImpl()
452455
return errors.New("connection handshake timeout")
453456
}
454457

455458
// Wait for initialization
456459
select {
457460
case <-w.initialized:
458-
// isInitialized is set by onEventFullStatus
461+
// isInitialized is set by onEventFullStatus
462+
w.logger.WithFields(logrus.Fields{"wattpilot": w.host, "initialized": w.isInitialized, "auhtenticated": w.isConnected}).Info("Initialization done")
463+
459464
case <-ctx.Done():
460-
w.disconnectImpl()
461-
return errors.New("initialization timeout")
465+
if !w.isInitialized {
466+
w.logger.WithFields(logrus.Fields{"wattpilot": w.host, "initialized": w.isInitialized, "auhtenticated": w.isConnected}).Info("Initialization timeout")
467+
w.disconnectImpl()
468+
return errors.New("initialization timeout")
469+
}
462470
}
463471

464472
return nil
@@ -492,7 +500,8 @@ func (w *Wattpilot) Disconnect() {
492500
}
493501

494502
func (w *Wattpilot) disconnectImpl() {
495-
w.logger.WithFields(logrus.Fields{"wattpilot": w.host}).Info("Closing connection...")
503+
504+
w.logger.WithFields(logrus.Fields{"wattpilot": w.host, "caller": getCallerFunctionName()}).Info("Closing connection...")
496505

497506
if w.conn == nil {
498507
return // Already disconnected
@@ -711,14 +720,16 @@ func (w *Wattpilot) onEventFullStatus(message map[string]interface{}) {
711720
if isPartial {
712721
return
713722
}
723+
724+
w.logger.WithFields(logrus.Fields{"wattpilot": w.host}).Trace("Initialization done")
725+
714726
if w.IsInitialized() {
715727
return
716728
}
717729

718-
w.logger.WithFields(logrus.Fields{"wattpilot": w.host}).Trace("Initialization done")
719-
720730
w.initialized <- true
721731
w.isInitialized = true
732+
722733
}
723734

724735
func (w *Wattpilot) onEventDeltaStatus(message map[string]interface{}) {

wattpilot_test.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ func TestConnect(t *testing.T) {
5555
t.Skip("WATTPILOT_HOST and WATTPILOT_PASSWORD environment variables not set. Skipping integration test.")
5656
}
5757

58-
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
58+
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
5959
defer cancel()
6060

6161
w := New(host, pwd)
@@ -70,7 +70,7 @@ func TestConnect(t *testing.T) {
7070
assert.NoError(t, err, "Connect should not return an error")
7171
assert.True(t, w.IsInitialized(), "Wattpilot should be initialized after successful connection")
7272
case <-ctx.Done():
73-
assert.Fail(t, "Test timed out after 60 seconds", ctx.Err())
73+
assert.Fail(t, "Test timed out after 30 seconds", ctx.Err())
7474
}
7575

7676
w.Disconnect()

0 commit comments

Comments
 (0)