From 0d39fb20110d2186f396ee81ce5e9f792ab0de95 Mon Sep 17 00:00:00 2001
From: Fred Lotter <fred.lotter@canonical.com>
Date: Mon, 3 Jul 2023 08:57:19 +0200
Subject: [PATCH] Review Changes #2

---
 internals/daemon/daemon_test.go | 20 +++++++++++++++++---
 1 file changed, 17 insertions(+), 3 deletions(-)

diff --git a/internals/daemon/daemon_test.go b/internals/daemon/daemon_test.go
index b2bb4293b..82a1c75a0 100644
--- a/internals/daemon/daemon_test.go
+++ b/internals/daemon/daemon_test.go
@@ -1175,9 +1175,14 @@ func (s *daemonSuite) TestSyscallPosRebootDelay(c *C) {
 	})()
 
 	period := time.Millisecond * 25
+	timeout := time.Second * 10
 	syscallReboot(period)
 	start := time.Now()
-	<-wait
+	select {
+	case <-wait:
+	case <-time.After(timeout): // exit test if we fail and get stuck
+		c.Fail()
+	}
 	elapse := time.Now().Sub(start)
 	c.Assert(elapse >= period, Equals, true)
 }
@@ -1198,7 +1203,11 @@ func (s *daemonSuite) TestSyscallNegRebootDelay(c *C) {
 	period := time.Second * 10
 	syscallReboot(-period)
 	start := time.Now()
-	<-wait
+	select {
+	case <-wait:
+	case <-time.After(period): // exit test if we fail and get stuck
+		c.Fail()
+	}
 	elapse := time.Now().Sub(start)
 	c.Assert(elapse < period, Equals, true)
 }
@@ -1222,7 +1231,12 @@ func (s *daemonSuite) TestSetSyscall(c *C) {
 	err := rebootHandler(0)
 	c.Assert(err, IsNil)
 	// This would block forever if the switch did not work.
-	<-wait
+	timeout := time.Second * 10
+	select {
+	case <-wait:
+	case <-time.After(timeout): // exit test if we fail and get stuck
+		c.Fail()
+	}
 }
 
 func (s *daemonSuite) TestCapSysBootFail(c *C) {