From 18d957853c579848d73bc087196985eaacb76b63 Mon Sep 17 00:00:00 2001
From: Ian Davis <18375+iand@users.noreply.github.com>
Date: Mon, 11 Sep 2023 09:50:09 +0100
Subject: [PATCH] fix: EventProbeRemove should also remove from routing table
 (#117)

---
 routing/probe.go      | 4 ++++
 routing/probe_test.go | 4 ++--
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/routing/probe.go b/routing/probe.go
index 572024d..cc755be 100644
--- a/routing/probe.go
+++ b/routing/probe.go
@@ -160,7 +160,11 @@ func (p *Probe[K, A]) Advance(ctx context.Context, ev ProbeEvent) ProbeState {
 		p.nvl.Put(nv)
 	case *EventProbeRemove[K]:
 		span.SetAttributes(attribute.String("event", "EventProbeRemove"), attribute.String("nodeid", tev.NodeID.String()))
+		p.rt.RemoveKey(tev.NodeID.Key())
 		p.nvl.Remove(tev.NodeID)
+		return &StateProbeNodeFailure[K]{
+			NodeID: tev.NodeID,
+		}
 	case *EventProbeMessageResponse[K, A]:
 		span.SetAttributes(attribute.String("event", "EventProbeMessageResponse"), attribute.String("nodeid", tev.NodeInfo.ID().String()))
 		nv, found := p.nvl.Get(tev.NodeInfo.ID())
diff --git a/routing/probe_test.go b/routing/probe_test.go
index a376809..4e2987d 100644
--- a/routing/probe_test.go
+++ b/routing/probe_test.go
@@ -281,8 +281,8 @@ func TestProbeRemoveDeletesNodeValue(t *testing.T) {
 		NodeID: kadtest.NewID(key.Key8(4)),
 	})
 
-	// state should remain idle
-	require.IsType(t, &StateProbeIdle{}, state)
+	// state indicate that node failed
+	require.IsType(t, &StateProbeNodeFailure[key.Key8]{}, state)
 
 	// advance time by one revisit interval
 	clk.Add(cfg.CheckInterval)