| import importlib.util,json,tempfile,unittest,sys,subprocess,os |
| from pathlib import Path |
| from unittest.mock import patch |
| P=Path(__file__).resolve().parents[1] |
| s=importlib.util.spec_from_file_location('inf',P/'inference.py');inf=importlib.util.module_from_spec(s);s.loader.exec_module(inf) |
| Q={'q_no':1,'question':'Example?','subject':'Demo','category':'stem','choices':['a','b','c','d'],'_level':'easy','_file':'easy/demo.json'} |
| class PipelineTests(unittest.TestCase): |
| def test_rotation_mapping_and_fallback(self): |
| k=inf.kof(Q) |
| with tempfile.TemporaryDirectory() as d: |
| with patch.object(inf,'chat',return_value=('no answer','')) as c: |
| result=inf.stage_rot([Q],['http://127.0.0.1:11434'],1,Path(d),1,'pass2','p2.jsonl',{k:'4'},{k:'3'}) |
| self.assertEqual(result[k],'2'); self.assertEqual(c.call_args.args[-1],300);self.assertEqual(c.call_count,1) |
| with patch.object(inf,'chat',return_value=('no answer','')): |
| self.assertEqual(inf.stage_rot([Q],['http://127.0.0.1:11434'],1,Path(d),2,'p4','p4.jsonl',{k:'4'},{k:'3'},False),{}) |
| def test_pass3_preserves_actual_sampling(self): |
| with tempfile.TemporaryDirectory() as d,patch.object(inf,'chat',return_value=('Answer: 2','')) as c: |
| inf.stage_pass3([Q],['http://127.0.0.1:11434'],1,Path(d),{inf.kof(Q):'4'}) |
| self.assertEqual(c.call_args.args[4],{'temperature':0.1,'top_p':0.95,'num_predict':4096}) |
| def test_transport_does_not_invent_seed(self): |
| from io import BytesIO |
| with patch.object(inf._OPENER,'open',return_value=BytesIO(b'{"message":{"content":"Answer: 2"}}')) as op: |
| inf.chat('http://127.0.0.1:11434','s','u',True,{'temperature':0.1}) |
| payload=json.loads(op.call_args.args[0].data) |
| self.assertNotIn('seed',payload['options']) |
| def test_loader_row_order_and_duplicates(self): |
| rows=[{**Q,'_file':'hard/b.json','_level':'hard'}, {**Q,'_file':'easy/a.json'}] |
| self.assertEqual([q['_file'] for q in inf.prepare(rows)],['hard/b.json','easy/a.json']) |
| with self.assertRaises(ValueError): inf.prepare([dict(Q),dict(Q)]) |
| def test_missing_or_duplicate_checkpoint_fails(self): |
| with tempfile.TemporaryDirectory() as d: |
| p=Path(d)/'p.jsonl' |
| with self.assertRaises(FileNotFoundError):inf.load_ck(p,required=True) |
| r={'q_no':1,'subject':'Demo','category':'stem','level':'easy','answer':'1'};p.write_text((json.dumps(r)+'\n')*2) |
| with self.assertRaises(ValueError):inf.load_ck(p) |
| def test_protect_submitted_file(self): |
| old=(P/'predictions.jsonl').read_bytes() |
| r=subprocess.run([sys.executable,'-B',str(P/'inference.py'),'--output',str(P/'predictions.jsonl')],capture_output=True,text=True) |
| self.assertNotEqual(r.returncode,0);self.assertIn('Do not overwrite',r.stderr);self.assertEqual(old,(P/'predictions.jsonl').read_bytes()) |
| def test_full_orchestration_with_scripted_responses(self): |
| |
| with tempfile.TemporaryDirectory() as td: |
| root=Path(td);data=root/'data' |
| for level in ['easy','medium','hard']: |
| (data/level).mkdir(parents=True) |
| (data/level/'demo.json').write_text(json.dumps([{'q_no':1,'subject':'Demo','category':'stem','question':'Example?','choices':['a','b','c','d']}])) |
| calls=[] |
| def chat(host,system,user,think,options,timeout=420): |
| calls.append(options) |
| if not think:return ('Answer: 1' if options['num_predict']==260 else '2','') |
| return ('Answer: 3','') |
| env={'WORKDIR':str(root/'ck'),'WORKERS':'1','OLLAMA_HOSTS':'http://127.0.0.1:11434'} |
| with patch.dict(os.environ,env),patch.object(sys,'argv',['inference.py','--local-data',str(data),'--output',str(root/'out.jsonl')]),patch.object(inf,'ensure_ollama'),patch.object(inf,'chat',side_effect=chat): |
| inf.main() |
| answers=[json.loads(l) for l in (root/'out.jsonl').read_text().splitlines()] |
| self.assertEqual(len(answers),3); self.assertTrue(any(x.get('temperature')==0.1 for x in calls)) |
| before=len(calls) |
| with patch.dict(os.environ,env),patch.object(sys,'argv',['inference.py','--local-data',str(data),'--output',str(root/'out2.jsonl')]),patch.object(inf,'ensure_ollama'),patch.object(inf,'chat',side_effect=chat):inf.main() |
| self.assertEqual(len(calls),before) |
| if __name__=='__main__':unittest.main(verbosity=2) |
|
|
