Clean up tests

2 files changed, 14 insertions, 6 deletions

diff --git a/src/tests/integration_test_utils.rs b/src/tests/integration_test_utils.rs
index 1395d97b..62ce041f 100644
--- a/src/tests/integration_test_utils.rs
+++ b/src/tests/integration_test_utils.rs
@@ -20,11 +20,13 @@ pub mod integration_test_utils {
 
     pub fn get_line_of_code_from_delta<'a>(
         input: &str,
+        line_number: usize,
+        expected_text: &str,
         options: cli::Opt,
     ) -> (String, config::Config<'a>) {
         let (output, config) = run_delta(&input, options);
-        let line_of_code = output.lines().nth(12).unwrap();
-        assert!(strip_ansi_codes(line_of_code) == " class X:");
+        let line_of_code = output.lines().nth(line_number).unwrap();
+        assert!(strip_ansi_codes(line_of_code) == expected_text);
         (line_of_code.to_string(), config)
     }
 
diff --git a/src/tests/test_example_diffs.rs b/src/tests/test_example_diffs.rs
index f46e2169..a01c94ce 100644
--- a/src/tests/test_example_diffs.rs
+++ b/src/tests/test_example_diffs.rs
@@ -47,8 +47,12 @@ mod tests {
     fn test_recognized_file_type() {
         // In addition to the background color, the code has language syntax highlighting.
         let options = integration_test_utils::get_command_line_options();
-        let (output, config) =
-            integration_test_utils::get_line_of_code_from_delta(&ADDED_FILE_INPUT, options);
+        let (output, config) = integration_test_utils::get_line_of_code_from_delta(
+            &ADDED_FILE_INPUT,
+            12,
+            " class X:",
+            options,
+        );
         ansi_test_utils::assert_has_color_other_than_plus_color(&output, &config);
     }
 
@@ -58,7 +62,8 @@ mod tests {
         // .txt syntax under the theme.
         let options = integration_test_utils::get_command_line_options();
         let input = ADDED_FILE_INPUT.replace("a.py", "a");
-        let (output, config) = integration_test_utils::get_line_of_code_from_delta(&input, options);
+        let (output, config) =
+            integration_test_utils::get_line_of_code_from_delta(&input, 12, " class X:", options);
         ansi_test_utils::assert_has_color_other_than_plus_color(&output, &config);
     }
 
@@ -69,7 +74,8 @@ mod tests {
         let mut options = integration_test_utils::get_command_line_options();
         options.theme = Some("none".to_string());
         let input = ADDED_FILE_INPUT.replace("a.py", "a");
-        let (output, config) = integration_test_utils::get_line_of_code_from_delta(&input, options);
+        let (output, config) =
+            integration_test_utils::get_line_of_code_from_delta(&input, 12, " class X:", options);
         ansi_test_utils::assert_has_plus_color_only(&output, &config);
     }