@@ -104,12 +104,13 @@ def test_main_success_exit_code(self, tmp_path):
104104 """
105105 # Mock run_agent to return success
106106 with patch ("src.main.run_agent" ) as mock_run :
107- mock_run .return_value = {
108- "job_status" : "completed" ,
109- "mode" : "proceed"
110- }
107+ with patch ("src.main._run_startup_preflight" , return_value = None ):
108+ mock_run .return_value = {
109+ "job_status" : "completed" ,
110+ "mode" : "proceed"
111+ }
111112
112- main (["--prompt" , "Test" , "--output-dir" , str (tmp_path )])
113+ main (["--prompt" , "Test" , "--output-dir" , str (tmp_path )])
113114
114115 def test_main_failure_exit_code (self , tmp_path ):
115116 """
@@ -120,16 +121,17 @@ def test_main_failure_exit_code(self, tmp_path):
120121 Then: sys.exit(1)
121122 """
122123 with patch ("src.main.run_agent" ) as mock_run :
123- mock_run .return_value = {
124- "job_status" : "failed" ,
125- "error" : "Test error" ,
126- "mode" : "fail"
127- }
124+ with patch ("src.main._run_startup_preflight" , return_value = None ):
125+ mock_run .return_value = {
126+ "job_status" : "failed" ,
127+ "error" : "Test error" ,
128+ "mode" : "fail"
129+ }
128130
129- with pytest .raises (SystemExit ) as exc_info :
130- main (["--prompt" , "Test" , "--output-dir" , str (tmp_path )])
131+ with pytest .raises (SystemExit ) as exc_info :
132+ main (["--prompt" , "Test" , "--output-dir" , str (tmp_path )])
131133
132- assert exc_info .value .code == 1
134+ assert exc_info .value .code == 1
133135
134136 def test_json_output_format (self , tmp_path , capsys ):
135137 """
@@ -142,19 +144,20 @@ def test_json_output_format(self, tmp_path, capsys):
142144 import json
143145
144146 with patch ("src.main.run_agent" ) as mock_run :
145- mock_run .return_value = {
146- "job_status" : "completed" ,
147- "job_id" : "test_123"
148- }
147+ with patch ("src.main._run_startup_preflight" , return_value = None ):
148+ mock_run .return_value = {
149+ "job_status" : "completed" ,
150+ "job_id" : "test_123"
151+ }
149152
150- main (["--prompt" , "Test" , "--output-dir" , str (tmp_path ), "--json" ])
153+ main (["--prompt" , "Test" , "--output-dir" , str (tmp_path ), "--json" ])
151154
152- captured = capsys .readouterr ()
155+ captured = capsys .readouterr ()
153156
154- # Verify valid JSON
155- output = json .loads (captured .out )
156- assert output ["job_status" ] == "completed"
157- assert output ["job_id" ] == "test_123"
157+ # Verify valid JSON
158+ output = json .loads (captured .out )
159+ assert output ["job_status" ] == "completed"
160+ assert output ["job_id" ] == "test_123"
158161
159162 def test_main_uses_short_dns_prompt_file (self , tmp_path ):
160163 """
@@ -169,13 +172,14 @@ def test_main_uses_short_dns_prompt_file(self, tmp_path):
169172 assert prompt_path .exists (), f"Prompt file missing: { prompt_path } "
170173
171174 with patch ("src.main.run_agent" ) as mock_run :
172- mock_run .return_value = {
173- "job_status" : "completed" ,
174- "mode" : "proceed"
175- }
175+ with patch ("src.main._run_startup_preflight" , return_value = None ):
176+ mock_run .return_value = {
177+ "job_status" : "completed" ,
178+ "mode" : "proceed"
179+ }
176180
177- main (["--prompt-path" , str (prompt_path ), "--output-dir" , str (tmp_path ), "--dry-run" ])
181+ main (["--prompt-path" , str (prompt_path ), "--output-dir" , str (tmp_path ), "--dry-run" ])
178182
179- assert mock_run .call_count == 1
180- passed_prompt = mock_run .call_args [0 ][0 ]
181- assert "50 steps" in passed_prompt
183+ assert mock_run .call_count == 1
184+ passed_prompt = mock_run .call_args [0 ][0 ]
185+ assert "50 steps" in passed_prompt
0 commit comments