From d0414dde9e6b99a656550132b881220b216eef53 Mon Sep 17 00:00:00 2001
From: Daniel D'Avella <dan.davella@pixee.ai>
Date: Fri, 17 Nov 2023 09:42:09 -0500
Subject: [PATCH] Update integration test

---
 integration_tests/base_test.py        |  3 +++
 integration_tests/test_url_sandbox.py | 21 +++++++++++++++++----
 tests/samples/make_request.py         |  4 +++-
 tests/samples/test_sources.py         |  2 ++
 4 files changed, 25 insertions(+), 5 deletions(-)
 create mode 100644 tests/samples/test_sources.py

diff --git a/integration_tests/base_test.py b/integration_tests/base_test.py
index 36808ebb..dd715c9c 100644
--- a/integration_tests/base_test.py
+++ b/integration_tests/base_test.py
@@ -4,12 +4,15 @@
 import os
 import pathlib
 import subprocess
+import sys
 
 from codemodder import __version__
 from codemodder import registry
 from tests.validations import execute_code
 
 SAMPLES_DIR = "tests/samples"
+# Enable import of test modules from test directory
+sys.path.append(SAMPLES_DIR)
 
 
 class CleanRepoMixin:
diff --git a/integration_tests/test_url_sandbox.py b/integration_tests/test_url_sandbox.py
index 288ef007..1fd9688a 100644
--- a/integration_tests/test_url_sandbox.py
+++ b/integration_tests/test_url_sandbox.py
@@ -11,13 +11,26 @@ class TestUrlSandbox(BaseIntegrationTest):
     original_code, expected_new_code = original_and_expected_from_code_path(
         code_path,
         [
-            (0, """from security import safe_requests\n"""),
-            (2, """safe_requests.get("https://www.google.com")\n"""),
+            (1, """from security import safe_requests\n"""),
+            (4, """safe_requests.get(url)\n"""),
         ],
     )
 
-    expected_diff = '--- \n+++ \n@@ -1,4 +1,4 @@\n-import requests\n+from security import safe_requests\n \n-requests.get("https://www.google.com")\n+safe_requests.get("https://www.google.com")\n var = "hello"\n'
-    expected_line_change = "3"
+    expected_diff = """\
+--- 
++++ 
+@@ -1,6 +1,6 @@
+ from test_sources import untrusted_data
+-import requests
++from security import safe_requests
+ 
+ url = untrusted_data()
+-requests.get(url)
++safe_requests.get(url)
+ var = "hello"
+"""
+
+    expected_line_change = "5"
     change_description = UrlSandbox.CHANGE_DESCRIPTION
     num_changed_files = 2
 
diff --git a/tests/samples/make_request.py b/tests/samples/make_request.py
index 0486bb89..f765edda 100644
--- a/tests/samples/make_request.py
+++ b/tests/samples/make_request.py
@@ -1,4 +1,6 @@
+from test_sources import untrusted_data
 import requests
 
-requests.get("https://www.google.com")
+url = untrusted_data()
+requests.get(url)
 var = "hello"
diff --git a/tests/samples/test_sources.py b/tests/samples/test_sources.py
new file mode 100644
index 00000000..1e53ad3b
--- /dev/null
+++ b/tests/samples/test_sources.py
@@ -0,0 +1,2 @@
+def untrusted_data():
+    return "http://www.google.com"