fix for tests

Author: mohamedashrraf222

codeflash/api/aiservice.py

@@ -32,7 +32,21 @@
 
 class AiServiceClient:
     def __init__(self) -> None:
-        # Validate API key before initializing the client
+        # API key validation is deferred to first use (in the headers property)
+        self.base_url = self.get_aiservice_base_url()
+        
+    def get_aiservice_base_url(self) -> str:
+        if os.environ.get("CODEFLASH_AIS_SERVER", default="prod").lower() == "local":
+            logger.info("Using local AI Service at http://localhost:8000")
+            console.rule()
+            return "http://localhost:8000"
+        return "https://app.codeflash.ai"
+    
+    def validate_api_key(self) -> None:
+        """Validate API key. Raises OSError if invalid."""
+        if hasattr(self, '_api_key_validated') and self._api_key_validated:
+            return
+            
         try:
             from codeflash.api.cfapi import get_user_id  # noqa: PLC0415
 
@@ -53,15 +67,14 @@ def __init__(self) -> None:
             # If cfapi is not available, skip validation
             pass
 
-        self.base_url = self.get_aiservice_base_url()
-        self.headers = {"Authorization": f"Bearer {get_codeflash_api_key()}", "Connection": "close"}
-
-    def get_aiservice_base_url(self) -> str:
-        if os.environ.get("CODEFLASH_AIS_SERVER", default="prod").lower() == "local":
-            logger.info("Using local AI Service at http://localhost:8000")
-            console.rule()
-            return "http://localhost:8000"
-        return "https://app.codeflash.ai"
+        self._api_key_validated = True
+    
+    @property
+    def headers(self) -> dict[str, str]:
+        """Get headers with API key. Validates API key on first access."""
+        # Lazily validate API key on first use
+        self.validate_api_key()
+        return {"Authorization": f"Bearer {get_codeflash_api_key()}", "Connection": "close"}
 
     def make_ai_service_request(
         self,

codeflash/optimization/function_optimizer.py

@@ -213,7 +213,7 @@ def __init__(
     ) -> None:
         self.project_root = test_cfg.project_root_path
         self.test_cfg = test_cfg
-        self.aiservice_client = aiservice_client if aiservice_client else AiServiceClient()
+        self._aiservice_client = aiservice_client  # Can be None for lazy initialization
         self.function_to_optimize = function_to_optimize
         self.function_to_optimize_source_code = (
             function_to_optimize_source_code
@@ -247,6 +247,13 @@ def __init__(
             max_workers=n_tests + 2 if self.experiment_id is None else n_tests + 3
         )
 
+    @property
+    def aiservice_client(self) -> AiServiceClient:
+        """Lazy initialization of AiServiceClient to delay API key validation."""
+        if self._aiservice_client is None:
+            self._aiservice_client = AiServiceClient()
+        return self._aiservice_client
+
     def can_be_optimized(self) -> Result[tuple[bool, CodeOptimizationContext, dict[Path, str]], str]:
         should_run_experiment = self.experiment_id is not None
         logger.debug(f"Function Trace ID: {self.function_trace_id}")

codeflash/optimization/optimizer.py

@@ -50,7 +50,7 @@ def __init__(self, args: Namespace) -> None:
             benchmark_tests_root=args.benchmarks_root if "benchmark" in args and "benchmarks_root" in args else None,
         )
 
-        self.aiservice_client = AiServiceClient()
+        self._aiservice_client: AiServiceClient | None = None
         self.experiment_id = os.getenv("CODEFLASH_EXPERIMENT_ID", None)
         self.local_aiservice_client = LocalAiServiceClient() if self.experiment_id else None
         self.replay_tests_dir = None
@@ -61,6 +61,13 @@ def __init__(self, args: Namespace) -> None:
         self.original_args_and_test_cfg: tuple[Namespace, TestConfig] | None = None
         self.patch_files: list[Path] = []
 
+    @property
+    def aiservice_client(self) -> AiServiceClient:
+        """Lazy initialization of AiServiceClient to delay API key validation."""
+        if self._aiservice_client is None:
+            self._aiservice_client = AiServiceClient()
+        return self._aiservice_client
+
     def run_benchmarks(
         self, file_to_funcs_to_optimize: dict[Path, list[FunctionToOptimize]], num_optimizable_functions: int
     ) -> tuple[dict[str, dict[BenchmarkKey, float]], dict[BenchmarkKey, float]]:
@@ -261,6 +268,10 @@ def run(self) -> None:
         console.rule()
         if not env_utils.ensure_codeflash_api_key():
             return
+        
+        # Validate API key before starting optimization
+        self.aiservice_client.validate_api_key()
+        
         if self.args.no_draft and is_pr_draft():
             logger.warning("PR is in draft mode, skipping optimization")
             return