From c2bdffb0abc789e80c8e119f4d59ccd6581fce0c Mon Sep 17 00:00:00 2001
From: Vladimir Zagainov <mrixs@mrixs.ru>
Date: Wed, 19 Aug 2026 13:26:05 +0300
Subject: [PATCH] fix: address code review findings

---
 internal/routing/graph.go      | 14 ++++++++++----
 internal/routing/graph_test.go | 19 +++++++++++++++----
 2 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/internal/routing/graph.go b/internal/routing/graph.go
index 9b1f8d6..70e01bd 100644
--- a/internal/routing/graph.go
+++ b/internal/routing/graph.go
@@ -1108,7 +1108,7 @@ func (g *Graph) checkRouteForChanges(itinerary *Itinerary) bool {
 				}
 				// Check for significant delay (more than 2x normal duration)
 				if edge.Duration > leg.Duration*2 && leg.Duration > 0 {
-					if !needsReSearch || itinerary.ReSearchReason == string(reasonNone) {
+					if !needsReSearch {
 						needsReSearch = true
 						itinerary.NeedsReSearch = true
 						itinerary.ReSearchReason = string(reasonMajorDelay)
@@ -1131,8 +1131,8 @@ func (g *Graph) rescheduleRoute(originID, destID string, opts SearchOptions, clo
 	result := g.FindRoute(originID, destID, opts, closedStations, neighbors, yclient...)
 	if result != nil {
 		result.LastChecked = time.Now().Unix()
-		result.NeedsReSearch = false
-		result.ReSearchReason = string(reasonNone)
+		// Keep NeedsReSearch and ReSearchReason from the original itinerary to indicate
+		// that a re-search was triggered due to changes
 	}
 	return result
 }
@@ -1141,7 +1141,13 @@ func (g *Graph) rescheduleRoute(originID, destID string, opts SearchOptions, clo
 // This is the main entry point for flight change notification logic.
 func (g *Graph) CheckAndRescheduleRoute(itinerary *Itinerary, originID, destID string, opts SearchOptions, closedStations map[string]bool, neighbors map[string][]storage.StationNeighbor, yclient ...*yandex.Client) *Itinerary {
 	if g.checkRouteForChanges(itinerary) {
-		return g.rescheduleRoute(originID, destID, opts, closedStations, neighbors, yclient...)
+		result := g.rescheduleRoute(originID, destID, opts, closedStations, neighbors, yclient...)
+		if result != nil {
+			// Preserve the NeedsReSearch and ReSearchReason from the original itinerary
+			result.NeedsReSearch = itinerary.NeedsReSearch
+			result.ReSearchReason = itinerary.ReSearchReason
+		}
+		return result
 	}
 	return itinerary
 }
diff --git a/internal/routing/graph_test.go b/internal/routing/graph_test.go
index e0cd091..4987060 100644
--- a/internal/routing/graph_test.go
+++ b/internal/routing/graph_test.go
@@ -413,6 +413,9 @@ func TestRouteReSearchOnChange(t *testing.T) {
 		}
 	}
 
+	// Reset LastChecked to force re-check (bypass the 1-hour cache)
+	itinerary.LastChecked = time.Now().Unix() - 7200
+
 	// Re-check for changes after simulating cancellation
 	checked2 := graph.CheckAndRescheduleRoute(itinerary, "s1", "s3", SearchOptions{MaxTransfers: 5}, nil, nil)
 	t.Logf("After cancellation - NeedsReSearch: %v, ReSearchReason: %s", checked2.NeedsReSearch, checked2.ReSearchReason)
@@ -440,12 +443,20 @@ func TestRouteReSearchOnChange(t *testing.T) {
 		ReSearchReason: "",
 	}
 
-	// For major delay, the check uses: edge.Duration > leg.Cost*2 && leg.Cost > 0
-	// With Cost=500, threshold would be 1000. Setting duration to 2000 should trigger.
+	// Reset the s1->s2 edge duration to normal value before testing major delay
+	for _, edge := range graph.edges {
+		if edge.From.ID == "s1" && edge.To.ID == "s2" {
+			edge.Duration = 3600 // Reset to normal duration
+			break
+		}
+	}
+
+	// For major delay, the check uses: edge.Duration > leg.Duration*2 && leg.Duration > 0
+	// With leg.Duration=3600, threshold would be 7200. Setting duration to 8000 should trigger.
 	for _, edge := range graph.edges {
 		if edge.From.ID == "s2" && edge.To.ID == "s3" {
-			edge.Duration = 2000 // > 500*2 = 1000, should trigger major delay
-			t.Logf("Set s2->s3 edge duration to %d (simulating major delay, threshold=1000)", edge.Duration)
+			edge.Duration = 8000 // > 3600*2 = 7200, should trigger major delay
+			t.Logf("Set s2->s3 edge duration to %d (simulating major delay, threshold=7200)", edge.Duration)
 			break
 		}
 	}