From 68e7978c8890fd0aec901e0871fb92061a0d8fa5 Mon Sep 17 00:00:00 2001
From: Matthew Farrellee <matt@cs.wisc.edu>
Date: Sat, 12 Jul 2025 19:53:54 -0400
Subject: [PATCH] chore: block network access from unit tests (#2732)

# What does this PR do?
this blocks network access for all `tests/unit/` tests.
`tests/integration/` are untouched.

it also introduces an `allow_network` marker to explicitly allow network
access.

## Test Plan
`./scripts/unit-tests.sh`
---
 pyproject.toml                                     |  4 ++++
 tests/unit/conftest.py                             | 11 +++++++++++
 tests/unit/providers/inference/test_remote_vllm.py |  1 +
 tests/unit/rag/test_vector_store.py                |  2 ++
 uv.lock                                            | 14 ++++++++++++++
 5 files changed, 32 insertions(+)

diff --git a/pyproject.toml b/pyproject.toml
index 9977d7372..2974ff996 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -64,6 +64,7 @@ dev = [
     "pytest-cov",
     "pytest-html",
     "pytest-json-report",
+    "pytest-socket",      # For blocking network access in unit tests
     "nbval",              # For notebook testing
     "black",
     "ruff",
@@ -344,3 +345,6 @@ classmethod-decorators = ["classmethod", "pydantic.field_validator"]
 
 [tool.pytest.ini_options]
 asyncio_mode = "auto"
+markers = [
+    "allow_network: Allow network access for specific unit tests",
+]
diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py
index aedac0386..b5eb1217d 100644
--- a/tests/unit/conftest.py
+++ b/tests/unit/conftest.py
@@ -4,6 +4,17 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
+import pytest_socket
+
 # We need to import the fixtures here so that pytest can find them
 # but ruff doesn't think they are used and removes the import. "noqa: F401" prevents them from being removed
 from .fixtures import cached_disk_dist_registry, disk_dist_registry, sqlite_kvstore  # noqa: F401
+
+
+def pytest_runtest_setup(item):
+    """Setup for each test - check if network access should be allowed."""
+    if "allow_network" in item.keywords:
+        pytest_socket.enable_socket()
+    else:
+        # Allowing Unix sockets is necessary for some tests that use local servers and mocks
+        pytest_socket.disable_socket(allow_unix_socket=True)
diff --git a/tests/unit/providers/inference/test_remote_vllm.py b/tests/unit/providers/inference/test_remote_vllm.py
index ca44cc95d..5c2ad03ab 100644
--- a/tests/unit/providers/inference/test_remote_vllm.py
+++ b/tests/unit/providers/inference/test_remote_vllm.py
@@ -393,6 +393,7 @@ async def test_process_vllm_chat_completion_stream_response_no_choices():
     assert chunks[0].event.event_type.value == "start"
 
 
+@pytest.mark.allow_network
 def test_chat_completion_doesnt_block_event_loop(caplog):
     loop = asyncio.new_event_loop()
     loop.set_debug(True)
diff --git a/tests/unit/rag/test_vector_store.py b/tests/unit/rag/test_vector_store.py
index dd36d3992..919f97ba7 100644
--- a/tests/unit/rag/test_vector_store.py
+++ b/tests/unit/rag/test_vector_store.py
@@ -123,6 +123,7 @@ class TestVectorStore:
         content = await content_from_doc(doc)
         assert content in DUMMY_PDF_TEXT_CHOICES
 
+    @pytest.mark.allow_network
     async def test_downloads_pdf_and_returns_content(self):
         # Using GitHub to host the PDF file
         url = "https://raw.githubusercontent.com/meta-llama/llama-stack/da035d69cfca915318eaf485770a467ca3c2a238/llama_stack/providers/tests/memory/fixtures/dummy.pdf"
@@ -135,6 +136,7 @@ class TestVectorStore:
         content = await content_from_doc(doc)
         assert content in DUMMY_PDF_TEXT_CHOICES
 
+    @pytest.mark.allow_network
     async def test_downloads_pdf_and_returns_content_with_url_object(self):
         # Using GitHub to host the PDF file
         url = "https://raw.githubusercontent.com/meta-llama/llama-stack/da035d69cfca915318eaf485770a467ca3c2a238/llama_stack/providers/tests/memory/fixtures/dummy.pdf"
diff --git a/uv.lock b/uv.lock
index bca12fc51..83e502e7f 100644
--- a/uv.lock
+++ b/uv.lock
@@ -1324,6 +1324,7 @@ dev = [
     { name = "pytest-cov" },
     { name = "pytest-html" },
     { name = "pytest-json-report" },
+    { name = "pytest-socket" },
     { name = "pytest-timeout" },
     { name = "ruamel-yaml" },
     { name = "ruff" },
@@ -1432,6 +1433,7 @@ dev = [
     { name = "pytest-cov" },
     { name = "pytest-html" },
     { name = "pytest-json-report" },
+    { name = "pytest-socket" },
     { name = "pytest-timeout" },
     { name = "ruamel-yaml" },
     { name = "ruff" },
@@ -2545,6 +2547,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/3e/43/7e7b2ec865caa92f67b8f0e9231a798d102724ca4c0e1f414316be1c1ef2/pytest_metadata-3.1.1-py3-none-any.whl", hash = "sha256:c8e0844db684ee1c798cfa38908d20d67d0463ecb6137c72e91f418558dd5f4b", size = 11428, upload-time = "2024-02-12T19:38:42.531Z" },
 ]
 
+[[package]]
+name = "pytest-socket"
+version = "0.7.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "pytest" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/05/ff/90c7e1e746baf3d62ce864c479fd53410b534818b9437413903596f81580/pytest_socket-0.7.0.tar.gz", hash = "sha256:71ab048cbbcb085c15a4423b73b619a8b35d6a307f46f78ea46be51b1b7e11b3", size = 12389, upload-time = "2024-01-28T20:17:23.177Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/19/58/5d14cb5cb59409e491ebe816c47bf81423cd03098ea92281336320ae5681/pytest_socket-0.7.0-py3-none-any.whl", hash = "sha256:7e0f4642177d55d317bbd58fc68c6bd9048d6eadb2d46a89307fa9221336ce45", size = 6754, upload-time = "2024-01-28T20:17:22.105Z" },
+]
+
 [[package]]
 name = "pytest-timeout"
 version = "2.4.0"