fix failing test

mohammedahmed18 · mohammedahmed18 · commit b5ca2b415e3a · 2025-12-11T19:00:23.000+02:00
diff --git a/tests/test_validate_python_code.py b/tests/test_validate_python_code.py
@@ -2,7 +2,7 @@
 from pydantic import ValidationError
 
 from codeflash.api.aiservice import AiServiceClient
-from codeflash.models.models import CodeString
+from codeflash.models.models import CodeString, OptimizedCandidateSource
 
 
 def test_python_string():
@@ -55,7 +55,7 @@ def test_generated_candidates_validation():
             "optimization_id": ""
         }
     ]
-    candidates = ai_service._get_valid_candidates(mock_generate_candidates)
+    candidates = ai_service._get_valid_candidates(mock_generate_candidates, OptimizedCandidateSource.OPTIMIZE)
     assert len(candidates) == 0
     code = """```python:file.py
 print('Hello, World!')
@@ -67,5 +67,5 @@ def test_generated_candidates_validation():
             "optimization_id": ""
         }
     ]
-    candidates = ai_service._get_valid_candidates(mock_generate_candidates)
+    candidates = ai_service._get_valid_candidates(mock_generate_candidates, OptimizedCandidateSource.OPTIMIZE)
     assert len(candidates) == 1

Original file line number	Diff line number	Diff line change
`@@ -2,7 +2,7 @@`
`2`	`2`	`from pydantic import ValidationError`
`3`	`3`
`4`	`4`	`from codeflash.api.aiservice import AiServiceClient`
`5`		`-from codeflash.models.models import CodeString`
	`5`	`+from codeflash.models.models import CodeString, OptimizedCandidateSource`
`6`	`6`
`7`	`7`
`8`	`8`	`def test_python_string():`
`@@ -55,7 +55,7 @@ def test_generated_candidates_validation():`
`55`	`55`	`"optimization_id": ""`
`56`	`56`	`}`
`57`	`57`	`]`
`58`		`- candidates = ai_service._get_valid_candidates(mock_generate_candidates)`
	`58`	`+ candidates = ai_service._get_valid_candidates(mock_generate_candidates, OptimizedCandidateSource.OPTIMIZE)`
`59`	`59`	`assert len(candidates) == 0`
`60`	`60`	code = """```python:file.py
`61`	`61`	`print('Hello, World!')`
`@@ -67,5 +67,5 @@ def test_generated_candidates_validation():`
`67`	`67`	`"optimization_id": ""`
`68`	`68`	`}`
`69`	`69`	`]`
`70`		`- candidates = ai_service._get_valid_candidates(mock_generate_candidates)`
	`70`	`+ candidates = ai_service._get_valid_candidates(mock_generate_candidates, OptimizedCandidateSource.OPTIMIZE)`
`71`	`71`	`assert len(candidates) == 1`