import tempfile,unittest,json
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import Mock,patch
from auto_test import AutoTests
from inference import Scheduler,InferenceError

class AutoTestsTests(unittest.TestCase):
 def setUp(self):
  self.tmp=tempfile.TemporaryDirectory();self.worker=Mock();self.worker.memory_plan={'gpu_layers':999};self.worker.catalog.entry.return_value=dict(id='m',kind='chat',file='model.gguf',profile_eligible=True)
  self.profiles=Mock();self.scheduler=Scheduler(self.worker);self.auto=AutoTests(Path(self.tmp.name)/'auto.json',self.profiles,self.worker,self.scheduler)
  self.gpus=[dict(name='RTX 5080',uuid='first'),dict(name='RTX 3060',uuid='second')]
 def tearDown(self):self.tmp.cleanup()
 def start(self):
  with patch('auto_test.probe',return_value=self.gpus):self.auto.start(dict(model_id='m',max_context=4096,mtp=False))
  self.auto.thread.join(3);self.assertFalse(self.auto.thread.is_alive())
 def test_primary_first_results_persist_and_save_explicitly(self):
  self.auto.request=Mock(return_value={});self.auto.benchmark=Mock(return_value={'predicted_per_second':40})
  self.start();job=self.auto.status()['job'];self.assertEqual(job['state'],'complete');self.assertEqual(len(job['results']),6)
  self.assertTrue(all(r['parameters']['gpu_devices']==['first'] for r in job['results']));self.profiles.save.assert_not_called();self.assertEqual(self.scheduler.active,0)
  self.auto.save(dict(job_id=job['id'],index=0,name='saved'));self.profiles.save.assert_called_once()
  self.assertEqual(AutoTests(self.auto.path,self.profiles,self.worker,self.scheduler).job['state'],'complete')
 def test_secondary_before_cpu(self):
  def plan(p,cancel):
   if p['parameters']['gpu_offload']=='full':raise InferenceError('does not fit')
  self.worker.plan.side_effect=plan;self.auto.request=Mock(return_value={});self.auto.benchmark=Mock(return_value={'predicted_per_second':10});self.start()
  rows=self.auto.job['results'];first_success=next(i for i,r in enumerate(rows) if r['success']);self.assertEqual(first_success,18);self.assertEqual(rows[first_success]['parameters']['gpu_offload'],'auto')
 def test_busy_lease_does_not_stop_another_worker(self):
  with self.scheduler.lease(('other',)):
   self.worker.reset_mock();self.start();self.assertEqual(self.auto.job['state'],'failed');self.worker.stop.assert_not_called()
 def test_cancel_releases_lease_and_does_not_save_profile(self):
  self.worker.plan.side_effect=lambda *a,**k:self.auto.cancel.set();self.auto.request=Mock(side_effect=InterruptedError());self.start();self.assertEqual(self.auto.job['state'],'cancelled');self.assertEqual(self.scheduler.active,0);self.profiles.save.assert_not_called()
 def test_validation_and_failed_result_save(self):
  with self.assertRaises(ValueError):self.auto.start(dict(model_id='m',max_context=True,mtp=False))
  self.auto.job={'id':'a','model_id':'m','results':[{'success':False}]}
  with self.assertRaises(ValueError):self.auto.save(dict(job_id='a',index=0,name='x'))
 def test_prompt_filled_by_token_count(self):
  self.auto.request=Mock(side_effect=[{'tokens':list(range(4000))},{'timings':{'prompt_n':3072,'predicted_n':64,'predicted_per_second':12}}]);r=self.auto.benchmark(4096);self.assertEqual(r['prompt_n'],3072);self.assertEqual(len(self.auto.request.call_args.args[1]['prompt']),3072)
if __name__=='__main__':unittest.main()