@@ -95,7 +95,7 @@ var _ = Describe("HealthMonitor", func() {
9595 Expect (fetched .Status ).To (Equal (StatusHealthy ))
9696 })
9797
98- It ("recovers unhealthy node when gRPC check succeeds " , func () {
98+ It ("recovers unhealthy node when heartbeat is fresh " , func () {
9999 node := makeNode ("unhealthy-worker" , "10.0.0.5:50051" , 8_000_000_000 )
100100 Expect (registry .Register (context .Background (), node , true )).To (Succeed ())
101101 Expect (node .Status ).To (Equal (StatusHealthy ))
@@ -106,15 +106,11 @@ var _ = Describe("HealthMonitor", func() {
106106 Expect (err ).ToNot (HaveOccurred ())
107107 Expect (fetched .Status ).To (Equal (StatusUnhealthy ))
108108
109- // Load a model so gRPC check is attempted
110- Expect (registry .SetNodeModel (context .Background (), node .ID , "test-model" , "loaded" , "10.0.0.5:50052" )).To (Succeed ())
111-
112- // Create health monitor with a factory that returns healthy clients
113- factory := newFakeBackendClientFactory ()
114- factory .setClient ("10.0.0.5:50051" , & fakeBackendClient {healthy : true })
115- hmWithFactory := NewHealthMonitor (registry , nil , 15 * time .Second , 30 * time .Second , "" , false , factory )
109+ // Refresh heartbeat (simulates the worker sending a heartbeat)
110+ Expect (db .Model (& BackendNode {}).Where ("id = ?" , node .ID ).
111+ Update ("last_heartbeat" , time .Now ()).Error ).ToNot (HaveOccurred ())
116112
117- hmWithFactory .doCheckAll (context .Background ())
113+ hm .doCheckAll (context .Background ())
118114
119115 fetched , err = registry .Get (context .Background (), node .ID )
120116 Expect (err ).ToNot (HaveOccurred ())
@@ -207,44 +203,77 @@ var _ = Describe("HealthMonitor (mock-based)", func() {
207203 Expect (calls ).NotTo (ContainElement (ContainSubstring ("MarkOffline" )))
208204 })
209205
210- It ("marks node unhealthy when gRPC check fails " , func () {
206+ It ("keeps node healthy when heartbeat is fresh (with models loaded) " , func () {
211207 store := newFakeNodeHealthStore ()
212208 factory := newFakeBackendClientFactory ()
213209 hm := newTestHealthMonitor (store , factory , true , staleThreshold )
214210
215- node := makeTestNode ("node-5" , "failing -worker" , "10.0.0.5:50051" , StatusHealthy , freshTime ())
211+ node := makeTestNode ("node-5" , "active -worker" , "10.0.0.5:50051" , StatusHealthy , freshTime ())
216212 store .addNode (node )
217213 store .addNodeModel ("node-5" , NodeModel {NodeID : "node-5" , ModelName : "llama-7b" })
218214
219- // Configure gRPC health check to fail
220- factory .setClient ("10.0.0.5:50051" , & fakeBackendClient {
221- healthy : false ,
222- err : fmt .Errorf ("connection refused" ),
223- })
224-
215+ // No gRPC client needed — health is determined by heartbeat, not gRPC probe
225216 hm .doCheckAll (context .Background ())
226217
227- Expect (store .getNode ("node-5" ).Status ).To (Equal (StatusUnhealthy ))
228- Expect (store .getCalls ()).To (ContainElement ("MarkUnhealthy:node-5" ))
218+ Expect (store .getNode ("node-5" ).Status ).To (Equal (StatusHealthy ))
219+ calls := store .getCalls ()
220+ Expect (calls ).NotTo (ContainElement (ContainSubstring ("MarkUnhealthy" )))
229221 })
230222
231- It ("recovers unhealthy node when gRPC check succeeds " , func () {
223+ It ("recovers unhealthy node when heartbeat is fresh " , func () {
232224 store := newFakeNodeHealthStore ()
233225 factory := newFakeBackendClientFactory ()
234226 hm := newTestHealthMonitor (store , factory , true , staleThreshold )
235227
236228 node := makeTestNode ("node-6" , "recovering-worker" , "10.0.0.6:50051" , StatusUnhealthy , freshTime ())
237229 store .addNode (node )
238- store .addNodeModel ("node-6" , NodeModel {NodeID : "node-6" , ModelName : "llama-7b" })
239-
240- // Configure gRPC health check to succeed
241- factory .setClient ("10.0.0.6:50051" , & fakeBackendClient {healthy : true })
242230
243231 hm .doCheckAll (context .Background ())
244232
245- // Should have called MarkHealthy to recover the node
246233 Expect (store .getCalls ()).To (ContainElement ("MarkHealthy:node-6" ))
247234 Expect (store .getNode ("node-6" ).Status ).To (Equal (StatusHealthy ))
248235 })
236+
237+ It ("node stays healthy when gRPC backend crashes but heartbeat is fresh" , func () {
238+ store := newFakeNodeHealthStore ()
239+ factory := newFakeBackendClientFactory ()
240+ hm := newTestHealthMonitor (store , factory , true , staleThreshold )
241+
242+ // Worker has a model loaded but the backend process crashed —
243+ // node should remain healthy because heartbeat is fresh
244+ node := makeTestNode ("node-crash" , "crash-worker" , "10.0.0.9:50051" , StatusHealthy , freshTime ())
245+ store .addNode (node )
246+ store .addNodeModel ("node-crash" , NodeModel {NodeID : "node-crash" , ModelName : "piper-model" , Address : "10.0.0.9:50053" })
247+
248+ // gRPC backend is dead — but health is heartbeat-based, not gRPC-based
249+ factory .setClient ("10.0.0.9:50051" , & fakeBackendClient {healthy : false , err : fmt .Errorf ("connection refused" )})
250+
251+ hm .doCheckAll (context .Background ())
252+
253+ Expect (store .getNode ("node-crash" ).Status ).To (Equal (StatusHealthy ))
254+ calls := store .getCalls ()
255+ Expect (calls ).NotTo (ContainElement (ContainSubstring ("MarkUnhealthy" )))
256+ })
257+
258+ It ("removes stale model via per-model health check without affecting node status" , func () {
259+ store := newFakeNodeHealthStore ()
260+ factory := newFakeBackendClientFactory ()
261+ hm := newTestHealthMonitor (store , factory , true , staleThreshold )
262+ hm .perModelHealthCheck = true
263+
264+ node := makeTestNode ("node-model" , "model-worker" , "10.0.0.10:50051" , StatusHealthy , freshTime ())
265+ store .addNode (node )
266+ store .addNodeModel ("node-model" , NodeModel {NodeID : "node-model" , ModelName : "piper-model" , Address : "10.0.0.10:50053" })
267+
268+ // Model backend is dead
269+ factory .setClient ("10.0.0.10:50053" , & fakeBackendClient {healthy : false , err : fmt .Errorf ("connection refused" )})
270+
271+ hm .doCheckAll (context .Background ())
272+
273+ // Node should remain healthy — only the model record is removed
274+ Expect (store .getNode ("node-model" ).Status ).To (Equal (StatusHealthy ))
275+ Expect (store .getCalls ()).To (ContainElement ("RemoveNodeModel:node-model:piper-model" ))
276+ Expect (store .getCalls ()).NotTo (ContainElement (ContainSubstring ("MarkUnhealthy" )))
277+ })
249278 })
250279})
0 commit comments