Refactor main

main.go

@@ -16,35 +16,54 @@ const (
 
 
 func main() {
-	var wg sync.WaitGroup
-	wg.Add(1)
-	stopFlag := atomicflag.New(false)
-	keypressFlag := atomicflag.New(false)
+	al := newAfterLock()
+	al.start()
+}
 
-	go func() {
-		for {
-			err := evdev.BlockUntilKeypress(keyboardDeviceNode)
-			if err != nil {
-				panic(err)
-			}
-			if stopFlag.Get() {
-				wg.Done()
-				return
-			} else {
-				keypressFlag.Set(true)
-			}
-		}
-	}()
+type afterLock struct {
+	wg sync.WaitGroup
+	stopFlag *atomicflag.AtomicFlag
+	keypressFlag *atomicflag.AtomicFlag
+}
 
+func newAfterLock() *afterLock {
+	return &afterLock{
+		stopFlag: atomicflag.New(false),
+		keypressFlag: atomicflag.New(false),
+	}
+}
+
+func (af *afterLock) start() {
+	af.wg.Add(1)
+
+	go af.detectKeypresses()
 	time.Sleep(initialDelay * time.Second)
+	af.hybernateDisplayLoop()
+}
+
+func (af *afterLock) detectKeypresses() {
+	for {
+		err := evdev.BlockUntilKeypress(keyboardDeviceNode)
+		if err != nil {
+			panic(err)
+		}
+		if af.stopFlag.Get() {
+			af.wg.Done()
+			return
+		}
+		af.keypressFlag.Set(true)
+	}
+}
+
+func (af *afterLock) hybernateDisplayLoop() {
 	for {
 		active, err := screensaver.IsActive()
 		if err != nil {
 			panic(err)
 		}
 		if !active {
-			stopFlag.Set(true)
-			wg.Wait()
+			af.stopFlag.Set(true)
+			af.wg.Wait()
 			return
 		}
 
@@ -59,8 +78,8 @@ func main() {
 			}
 			for {
 				time.Sleep(loopDelay * time.Second)
-				if keypressFlag.Get() {
-					keypressFlag.Set(false)
+				if af.keypressFlag.Get() {
+					af.keypressFlag.Set(false)
 					continue
 				} else {
 					break