import tempfile
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch,Mock
from inference import LlamaWorker,InferenceError
from profiles import SCHEMAS,CHAT_GPU_DEFAULTS

class WorkerTests(unittest.TestCase):
 def setUp(self):
  self.tmp=tempfile.TemporaryDirectory();self.root=Path(self.tmp.name);directory=self.root/'runtime'/'build-id'/'build/bin';directory.mkdir(parents=True);(directory/'llama-server').touch()
  self.entry=dict(id='model',file='test.gguf',size=1024**3)
  self.runtime=SimpleNamespace(root=self.root/'runtime',status=lambda:dict(active='build-id',builds=[dict(id='build-id',backend='CUDA',fit_tool=True)]))
  self.worker=LlamaWorker(self.root/'worker',SimpleNamespace(root=self.root/'models',entry=lambda _:self.entry),self.runtime)
  self.profile=dict(id='p',name='medium',revision=1,model_id='model',parameters={**{k:v[2] for k,v in SCHEMAS['chat'].items()},**CHAT_GPU_DEFAULTS,'gpu_devices':['gpu-first','gpu-second'],'split_mode':'layer','tensor_split':[85,15]})
  self.gpus=[dict(uuid='gpu-first',processes=0,total_mib=16000,free_mib=15000),dict(uuid='gpu-second',processes=0,total_mib=12000,free_mib=11000)]
 def tearDown(self):self.tmp.cleanup()
 def patches(self):
  real=Path.read_text
  return patch.object(Path,'read_text',lambda path,*a,**kw:'MemAvailable: 50000000 kB\n' if str(path)=='/proc/meminfo' else real(path,*a,**kw))
 def test_fixed_split_uses_prediction_and_preserves_device_order(self):
  process=Mock();process.poll.return_value=None
  http=Mock();http.getresponse.return_value.status=200
  with self.patches(),patch('inference.probe',return_value=self.gpus),patch('inference.cgroup_headroom',return_value=32*1024**3),patch.object(self.worker,'_fit_command',return_value='CUDA0 9000 1000 1000\nCUDA1 2000 1000 1000\nHost 512 0 256\n') as fit,patch('inference.subprocess.Popen',return_value=process) as launch,patch('inference.http.client.HTTPConnection',return_value=http),patch('inference.threading.Thread'):
   self.worker.ensure(self.profile)
  args=launch.call_args.args[0];env=launch.call_args.kwargs['env']
  self.assertEqual(env['CUDA_VISIBLE_DEVICES'],'gpu-first,gpu-second');self.assertEqual(args[args.index('--tensor-split')+1],'85,15');self.assertIn('--kv-unified',args);self.assertIn('--fit-print',fit.call_args.args[0]);self.assertEqual(self.worker.status()['state'],'ready')
 def test_mtp_requires_adapter_before_loading(self):
  self.profile['parameters']['mtp']=True
  with patch('inference.subprocess.Popen') as launch:
   with self.assertRaisesRegex(InferenceError,'MTP'):self.worker.ensure(self.profile)
   launch.assert_not_called()
 def test_mtp_is_passed_to_fit_and_server(self):
  self.profile['parameters'].update(mtp=True,mtp_tokens=2,mtp_min_p=.05)
  (self.root/'runtime/build-id/build/bin/deck-mtp-fit-v1').touch()
  process=Mock();process.poll.return_value=None;http=Mock();http.getresponse.return_value.status=200
  with self.patches(),patch('inference.probe',return_value=self.gpus),patch('inference.cgroup_headroom',return_value=32*1024**3),patch.object(self.worker,'_fit_command',return_value='CUDA0 9000 1000 1000\nCUDA1 2000 1000 1000\nHost 512 0 256\n') as fit,patch('inference.subprocess.Popen',return_value=process) as launch,patch('inference.http.client.HTTPConnection',return_value=http),patch('inference.threading.Thread'):
   self.worker.ensure(self.profile)
  args=launch.call_args.args[0];self.assertEqual(args[args.index('--spec-type')+1],'draft-mtp');self.assertEqual(args[args.index('--spec-draft-n-max')+1],'2');self.assertIn('--deck-mtp',fit.call_args.args[0]);self.assertNotIn('--deck-mtp',args)
 def test_busy_gpu_is_rejected_before_fit_or_spawn(self):
  self.gpus[0]['processes']=1
  with patch('inference.probe',return_value=self.gpus),patch('inference.subprocess.Popen') as launch:
   with self.assertRaises(InferenceError):self.worker.ensure(self.profile)
   launch.assert_not_called()
 def test_fixed_split_exhaustion_does_not_load_model(self):
  with self.patches(),patch('inference.probe',return_value=self.gpus),patch('inference.cgroup_headroom',return_value=32*1024**3),patch.object(self.worker,'_fit_command',return_value='CUDA0 50000 1000 1000\nCUDA1 2000 1000 1000\nHost 512 0 256\n') as fit,patch('inference.subprocess.Popen') as launch:
   with self.assertRaises(InferenceError):self.worker.ensure(self.profile)
   launch.assert_not_called();self.assertLessEqual(fit.call_count,11)
 def test_build_symlink_cannot_escape_deck_state(self):
  self.runtime.status=lambda:dict(active='outside',builds=[dict(id='outside',backend='CUDA',fit_tool=True)])
  (self.root/'runtime'/'outside').symlink_to('/usr')
  with self.assertRaises(InferenceError):self.worker.build()

if __name__=='__main__':unittest.main()