From 838c7a50ec75b9e64a95b206f541863400d49668 Mon Sep 17 00:00:00 2001
From: Dylan Huang <dhuang@fireworks.ai>
Date: Tue, 13 Jan 2026 16:25:00 -0800
Subject: [PATCH] Refactor dotenv handling in auth module and integrate
 environment variable loading into local test command. Introduced functions to
 find and retrieve values from .env files, enhancing configuration management
 for Docker tests.

---
 eval_protocol/auth.py                    | 65 +++++++++++++++++++++++-
 eval_protocol/cli_commands/local_test.py |  7 +++
 2 files changed, 71 insertions(+), 1 deletion(-)

diff --git a/eval_protocol/auth.py b/eval_protocol/auth.py
index 68ce134c..19b3e76d 100644
--- a/eval_protocol/auth.py
+++ b/eval_protocol/auth.py
@@ -1,12 +1,75 @@
 import logging
 import os
-from typing import Optional
+from typing import Dict, Optional
 
 import requests
+from dotenv import dotenv_values, find_dotenv, load_dotenv
 
 logger = logging.getLogger(__name__)
 
 
+def find_dotenv_path(search_path: Optional[str] = None) -> Optional[str]:
+    """
+    Find the .env file path, searching .env.dev first, then .env.
+
+    Args:
+        search_path: Directory to search from. If None, uses current working directory.
+
+    Returns:
+        Path to the .env file if found, otherwise None.
+    """
+    # If a specific search path is provided, look there first
+    if search_path:
+        env_dev_path = os.path.join(search_path, ".env.dev")
+        if os.path.isfile(env_dev_path):
+            return env_dev_path
+        env_path = os.path.join(search_path, ".env")
+        if os.path.isfile(env_path):
+            return env_path
+        return None
+
+    # Otherwise use find_dotenv to search up the directory tree
+    env_dev_path = find_dotenv(filename=".env.dev", raise_error_if_not_found=False, usecwd=True)
+    if env_dev_path:
+        return env_dev_path
+    env_path = find_dotenv(filename=".env", raise_error_if_not_found=False, usecwd=True)
+    if env_path:
+        return env_path
+    return None
+
+
+def get_dotenv_values(search_path: Optional[str] = None) -> Dict[str, Optional[str]]:
+    """
+    Get all key-value pairs from the .env file.
+
+    Args:
+        search_path: Directory to search from. If None, uses current working directory.
+
+    Returns:
+        Dictionary of environment variable names to values.
+    """
+    dotenv_path = find_dotenv_path(search_path)
+    if dotenv_path:
+        return dotenv_values(dotenv_path)
+    return {}
+
+
+# --- Load .env files ---
+# Attempt to load .env.dev first, then .env as a fallback.
+# This happens when the module is imported.
+# We use override=False (default) so that existing environment variables
+# (e.g., set in the shell) are NOT overridden by .env files.
+_DOTENV_PATH = find_dotenv_path()
+if _DOTENV_PATH:
+    load_dotenv(dotenv_path=_DOTENV_PATH, override=False)
+    logger.debug(f"eval_protocol.auth: Loaded environment variables from: {_DOTENV_PATH}")
+else:
+    logger.debug(
+        "eval_protocol.auth: No .env.dev or .env file found. Relying on shell/existing environment variables."
+    )
+# --- End .env loading ---
+
+
 def get_fireworks_api_key() -> Optional[str]:
     """
     Retrieves the Fireworks API key.
diff --git a/eval_protocol/cli_commands/local_test.py b/eval_protocol/cli_commands/local_test.py
index 97e02e9f..43a59a3f 100644
--- a/eval_protocol/cli_commands/local_test.py
+++ b/eval_protocol/cli_commands/local_test.py
@@ -5,6 +5,7 @@
 import sys
 from typing import List
 
+from ..auth import get_dotenv_values
 from .utils import _build_entry_point, _discover_and_select_tests
 
 
@@ -71,6 +72,12 @@ def _run_pytest_in_docker(
         workdir,
     ]
 
+    # Forward environment variables from .env file to the container
+    dotenv_vars = get_dotenv_values(project_root)
+    for key, value in dotenv_vars.items():
+        if value is not None:
+            cmd += ["-e", f"{key}={value}"]
+
     # If EP_SUMMARY_JSON is set on the host, mirror it into the container so that
     # pytest evaluation tests can write summary artifacts that are visible to the
     # host. We map paths under the host logs directory (~/.eval_protocol) into the