import tempfile,unittest,json,time
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import Mock,patch
from auto_test import AutoTests,model_layers,fine_splits,MAX_CANDIDATES,GPU_SPLITS
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',start_context=4096,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']),3*(1+len(GPU_SPLITS)))
  self.assertEqual(job['completed_contexts'],[4096]);self.assertEqual(job['results'][0]['parameters']['gpu_devices'],['first']);self.profiles.save.assert_not_called();self.assertEqual(self.scheduler.active,0)
  tested_splits={tuple(r['parameters']['tensor_split']) for r in job['results']}
  self.assertIn((75,25),tested_splits);self.assertIn((50,50),tested_splits)
  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_kv_cache_choice_is_used_and_saved(self):
  self.auto.request=Mock(return_value={});self.auto.benchmark=Mock(return_value={'predicted_per_second':40})
  with patch('auto_test.probe',return_value=self.gpus):
   self.auto.start(dict(model_id='m',start_context=4096,max_context=4096,mtp=False,cache_type_k='q8_0',cache_type_v='f16'))
  self.auto.thread.join(3)
  row=next(r for r in self.auto.job['results'] if r['success'])
  self.assertEqual((row['parameters']['cache_type_k'],row['parameters']['cache_type_v']),('q8_0','f16'))
  self.auto.save(dict(job_id=self.auto.job['id'],index=self.auto.job['results'].index(row),name='KV Test'))
  self.assertEqual(self.profiles.save.call_args.args[0]['parameters']['cache_type_v'],'f16')
 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,3*(1+len(GPU_SPLITS)));self.assertEqual(rows[first_success]['parameters']['gpu_offload'],'auto')
  self.assertEqual(rows[0]['failure_stage'],'prediction')
 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))
  with self.assertRaises(ValueError):self.auto.start(dict(model_id='m',start_context=65536,max_context=8192,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':3968,'predicted_n':64,'predicted_per_second':12}}]);r=self.auto.benchmark(4096);self.assertEqual(r['prompt_n'],3968);self.assertEqual(len(self.auto.request.call_args.args[1]['prompt']),3968)
 def test_budget_exhaustion_is_not_reported_as_complete(self):
  self.auto.job=dict(id='j',model_id='m',start_context=65536,max_context=65536,mtp=False,attempts=MAX_CANDIDATES,results=[],completed_contexts=[],started_at=time.time(),max_candidates=MAX_CANDIDATES)
  self.auto.run('first','second')
  self.assertEqual(self.auto.job['state'],'partial');self.assertEqual(self.auto.job['completed_contexts'],[])
class FineTests(unittest.TestCase):
 def test_single_layer_steps_between_coarse_candidates(self):
  import math
  values=fine_splits(66,[85,15],90)
  self.assertEqual([k for k,_ in values],[58,59])
  for k,split in values:self.assertEqual(math.ceil(66*split[0]/sum(split)),k)
  self.assertEqual(fine_splits(66,[95,5],100)[-1][0],65)
 def test_gguf_metadata_mtp_and_output(self):
  import struct
  def text(s):b=s.encode();return struct.pack('<Q',len(b))+b
  blob=b'GGUF'+struct.pack('<IQQ',3,0,3)+text('general.architecture')+struct.pack('<I',8)+text('qwen35')
  blob+=text('qwen35.block_count')+struct.pack('<II',4,65)+text('qwen35.nextn_predict_layers')+struct.pack('<II',4,1)
  with tempfile.TemporaryDirectory() as d:
   p=Path(d)/'model.gguf';p.write_bytes(blob);self.assertEqual(model_layers(p,True),66);self.assertEqual(model_layers(p,False),65)
   p.write_bytes(blob[:-2])
   with self.assertRaises(ValueError):model_layers(p,True)
 def test_refinement_integrated_after_first_fit(self):
  fixture=AutoTestsTests();fixture.setUp()
  try:
   a=fixture.auto;a.job=dict(id='j',model_id='m',start_context=4096,max_context=4096,mtp=True,attempts=0,results=[],completed_contexts=[],started_at=time.time())
   seen=[]
   def candidate(context,tier,devices,split,ratio,offload,layer_count=None):
    seen.append((context,ratio,layer_count));return [True] if ratio and ratio[0]<=90 else []
   a.test_candidate=candidate;fixture.worker.catalog.root=Path('/unused')
   with patch('auto_test.model_layers',return_value=66):a.run('first','second')
   self.assertEqual([v[2] for v in seen if v[2] is not None],[61])
  finally:fixture.tearDown()
if __name__=='__main__':unittest.main()