47 lines
2.9 KiB
Python
47 lines
2.9 KiB
Python
import unittest
|
|
from comfy_progress import _hp_event, _hp_snapshot, _hp_history_estimate
|
|
|
|
class ProgressTests(unittest.TestCase):
|
|
def state(self):
|
|
return {'prompt':'one','classes':{'12':'SamplerCustomAdvanced','13':'VAEDecode'},'samples':[], 'connected':True}
|
|
def event(self,state,value,t,node='12',prompt='one',maximum=20):
|
|
_hp_event(state,{'type':'progress','data':{'prompt_id':prompt,'node':node,'value':value,'max':maximum}},t)
|
|
def test_job_isolation_and_node_reset(self):
|
|
s=self.state();self.event(s,1,100);self.event(s,2,110,prompt='other')
|
|
self.assertEqual(s['value'],1)
|
|
self.event(s,2,110);self.event(s,3,120)
|
|
self.assertEqual(_hp_snapshot(s,120)['samplingMinSeconds'],136)
|
|
_hp_event(s,{'type':'executing','data':{'prompt_id':'one','node':'13'}},125)
|
|
r=_hp_snapshot(s,125);self.assertEqual(r['phase'],'decoding');self.assertNotIn('step',r)
|
|
def test_stale_disconnect_reset_no_false_countdown(self):
|
|
s=self.state()
|
|
for i in range(1,4):self.event(s,i,100+i*10)
|
|
self.assertNotIn('samplingMinSeconds',_hp_snapshot(s,200))
|
|
self.event(s,1,201)
|
|
self.assertNotIn('samplingMinSeconds',_hp_snapshot(s,201))
|
|
s['connected']=False
|
|
self.assertEqual(_hp_snapshot(s,202),{'phase':'running'})
|
|
def test_sampling_completion_is_not_video_completion(self):
|
|
s=self.state()
|
|
for i in [18,19,20]:self.event(s,i,100+i*10)
|
|
r=_hp_snapshot(s,300)
|
|
self.assertEqual(r['phase'],'sampling');self.assertEqual(r['step'],20)
|
|
self.assertNotIn('samplingMinSeconds',r)
|
|
def test_history_exact_settings_and_cache_exclusion(self):
|
|
graph={'1':{'class_type':'SamplerCustomAdvanced','inputs':{'steps':20,'noise_seed':123}},
|
|
'2':{'class_type':'MiniMaxH3ReferenceToVideo','inputs':{'width':1344,'height':768,'text':'private'}}}
|
|
def history(duration=1000,change=None,cached=False,failed=False):
|
|
import copy
|
|
other=copy.deepcopy(graph);other['1']['inputs']['noise_seed']=456;other['2']['inputs']['text']='another'
|
|
if change:other['2']['inputs'].update(change)
|
|
return {'prompt':[1,'id',other], 'status':{'status_str':'error' if failed else 'success','completed':True,
|
|
'messages':[['execution_start',{'timestamp':1000}],['execution_cached',{'nodes':['1'] if cached else []}],
|
|
['execution_success',{'timestamp':1000+duration*1000}]]}}
|
|
r=_hp_history_estimate(graph,{'a':history()})
|
|
self.assertEqual(r,{'historySamples':1})
|
|
r=_hp_history_estimate(graph,{'a':history(),'b':history(1200),'c':history(5,cached=True),
|
|
'd':history(500,change={'width':2560}),'e':history(2,failed=True)})
|
|
self.assertEqual(r,{'historySamples':2,'historyMinSeconds':800,'historyMaxSeconds':1440})
|
|
|
|
if __name__=='__main__':unittest.main()
|