From c24acb095c7efc335bf718c154569f15ad8d0596 Mon Sep 17 00:00:00 2001
From: Soutaro Matsumoto <soutaro.matsumoto@shopify.com>
Date: Thu, 7 May 2026 16:40:32 +0900
Subject: [PATCH] Add source code and test helper in GraphTest

`GraphTest::source_of(definition_id)` returns the source code for the definition.
---
 rust/rubydex/src/test_utils/graph_test.rs | 35 +++++++++++++++++++++--
 1 file changed, 33 insertions(+), 2 deletions(-)

diff --git a/rust/rubydex/src/test_utils/graph_test.rs b/rust/rubydex/src/test_utils/graph_test.rs
index 6a1a673c..c2ae33ec 100644
--- a/rust/rubydex/src/test_utils/graph_test.rs
+++ b/rust/rubydex/src/test_utils/graph_test.rs
@@ -1,20 +1,26 @@
+use std::collections::HashMap;
+
 use super::normalize_indentation;
 #[cfg(test)]
 use crate::diagnostic::Rule;
 use crate::indexing::{self, LanguageId};
 use crate::model::graph::{Graph, NameDependent};
-use crate::model::ids::{NameId, StringId};
+use crate::model::ids::{DefinitionId, NameId, StringId};
 use crate::resolution::Resolver;
 
 #[derive(Default)]
 pub struct GraphTest {
     graph: Graph,
+    sources: HashMap<String, String>,
 }
 
 impl GraphTest {
     #[must_use]
     pub fn new() -> Self {
-        Self { graph: Graph::new() }
+        Self {
+            graph: Graph::new(),
+            sources: HashMap::new(),
+        }
     }
 
     #[must_use]
@@ -31,12 +37,37 @@ impl GraphTest {
     pub fn index_uri(&mut self, uri: &str, source: &str) {
         let source = normalize_indentation(source);
         indexing::index_source(&mut self.graph, uri, &source, &LanguageId::Ruby);
+        self.sources.insert(uri.to_string(), source);
     }
 
     /// Indexes an RBS source
     pub fn index_rbs_uri(&mut self, uri: &str, source: &str) {
         let source = normalize_indentation(source);
         indexing::index_source(&mut self.graph, uri, &source, &LanguageId::Rbs);
+        self.sources.insert(uri.to_string(), source);
+    }
+
+    /// Returns the normalized source for the given URI.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the URI has not been indexed.
+    #[must_use]
+    pub fn source(&self, uri: &str) -> &str {
+        self.sources.get(uri).expect("source not found for URI")
+    }
+
+    /// Returns the source text for a definition, sliced by its offset.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the definition or its document does not exist.
+    #[must_use]
+    pub fn source_of(&self, definition_id: &DefinitionId) -> &str {
+        let def = self.graph.definitions().get(definition_id).unwrap();
+        let uri = self.graph.documents().get(def.uri_id()).unwrap().uri();
+        let source = self.source(uri);
+        &source[def.offset().start() as usize..def.offset().end() as usize]
     }
 
     pub fn delete_uri(&mut self, uri: &str) {