From 3e2050bf7e324c7828f8a7c31c2a2e146a3374a6 Mon Sep 17 00:00:00 2001
From: Vincent Raymond <vincent@lum.ai>
Date: Mon, 11 Mar 2024 12:06:01 -0400
Subject: [PATCH] Adding tests for module_locate

---
 pyproject.toml                                |  2 +-
 skema/program_analysis/module_locate.py       |  8 +++-
 .../tests/test_module_locate.py               | 41 +++++++++++++++++++
 3 files changed, 49 insertions(+), 2 deletions(-)
 create mode 100644 skema/program_analysis/tests/test_module_locate.py

diff --git a/pyproject.toml b/pyproject.toml
index 20028ab88d8..895eadbbc1d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -43,7 +43,7 @@ dynamic = ["readme"]
 #   Pygraphviz is often tricky to install, so we reserve it for the dev extras
 #   list.
 # - six: Required by auto-generated Swagger models
-dev = ["pytest", "pytest-cov", "pytest-xdist", "pytest-asyncio", "black", "mypy", "coverage", "pygraphviz", "six"]
+dev = ["pytest", "pytest-cov", "pytest-xdist", "pytest-asyncio", "pytest-mock", "black", "mypy", "coverage", "pygraphviz", "six"]
 
 demo = ["notebook"]
 
diff --git a/skema/program_analysis/module_locate.py b/skema/program_analysis/module_locate.py
index c9a43e9919d..2a2201a121c 100644
--- a/skema/program_analysis/module_locate.py
+++ b/skema/program_analysis/module_locate.py
@@ -16,7 +16,9 @@
 def identify_source_type(source: str):
     if not source:
         return "Unknown"
-    if "github" in source:
+    elif source == "https://github.com/python/cpython":
+        return "Compiled"
+    elif "github" in source:
         return "Repository"
     elif source.startswith("http"):
         return "Url"
@@ -53,6 +55,10 @@ def module_locate(module_name: str) -> str:
     :return: The module's file path, GitHub URL, or tarball URL.
     """
 
+    # Check if module is compiled into Python
+    if module_name in sys.builtin_module_names:
+        return "https://github.com/python/cpython"
+
     # Attempt to find the module in the local environment
     try:
         module_obj = importlib.import_module(module_name)
diff --git a/skema/program_analysis/tests/test_module_locate.py b/skema/program_analysis/tests/test_module_locate.py
new file mode 100644
index 00000000000..c95b4034ecc
--- /dev/null
+++ b/skema/program_analysis/tests/test_module_locate.py
@@ -0,0 +1,41 @@
+import pytest
+from unittest.mock import patch
+from skema.program_analysis.module_locate import identify_source_type, module_locate  
+
+# Testing identify_source_type
+@pytest.mark.parametrize("source,expected_type", [
+    ("https://github.com/python/cpython", "Compiled"),
+    ("https://github.com/other/repository", "Repository"),
+    ("http://example.com", "Url"),
+    ("local/path/to/module", "Local"),
+    ("", "Unknown"),
+])
+def test_identify_source_type(source, expected_type):
+    assert identify_source_type(source) == expected_type
+
+# Mocking requests.get to test module_locate without actual HTTP requests
+@pytest.fixture
+def mock_requests_get(mocker):
+    mock = mocker.patch('skema.program_analysis.module_locate.requests.get')
+    return mock
+
+def test_module_locate_builtin_module():
+    assert module_locate("sys") == "https://github.com/python/cpython"
+
+def test_module_locate_from_pypi_with_github_source(mock_requests_get):
+    mock_requests_get.return_value.json.return_value = {
+        'info': {'version': '1.0.0', 'project_urls': {'Source': 'https://github.com/example/project'}},
+        'releases': {'1.0.0': [{'filename': 'example-1.0.0.tar.gz', 'url': 'https://example.com/example-1.0.0.tar.gz'}]}
+    }
+    assert module_locate("example") == "https://github.com/example/project"
+
+def test_module_locate_from_pypi_with_tarball_url(mock_requests_get):
+    mock_requests_get.return_value.json.return_value = {
+        'info': {'version': '1.2.3'},
+        'releases': {'1.2.3': [{'filename': 'package-1.2.3.tar.gz', 'url': 'https://pypi.org/package-1.2.3.tar.gz'}]}
+    }
+    assert module_locate("package") == "https://pypi.org/package-1.2.3.tar.gz"
+
+def test_module_locate_not_found(mock_requests_get):
+    mock_requests_get.side_effect = Exception("Module not found")
+    assert module_locate("nonexistent") is None