yaoyao/tests/server/imageManagement.test.ts

// @vitest-environment node
import {it,expect,vi} from 'vitest'
import {DatabaseSync} from 'node:sqlite'
import {mkdtemp,rm,writeFile} from 'node:fs/promises'
import {tmpdir} from 'node:os'
import {join} from 'node:path'
import {randomUUID} from 'node:crypto'
import {ComputerRuntime} from '../../src/runner/worker/gateway'
import {ImageManagement} from '../../src/runner/computers/imageManagement'
import {fileSHA256} from '../../src/runner/computers/images'
it('persists image selection, rejects in-use switching and recovers interrupted jobs without replay',async()=>{
 const home=await mkdtemp(join(tmpdir(),'yaoyao-images-ui-')),db=new DatabaseSync(':memory:'),initial='sha256:'+'a'.repeat(64),next='sha256:'+'b'.repeat(64)
 const config:any={protocol:1,runnerId:randomUUID(),serverURL:'http://127.0.0.1',hermesURL:'http://127.0.0.1',token:'fixture',allowedProfiles:['default'],artifactRoots:[],computers:{runtime:'docker',imageId:initial,python:'/fixture/python',hermesSource:'/fixture/source',hermesHome:'/fixture/home'}}
 const runtime=new ComputerRuntime(db,config,home);let manager:ImageManagement|undefined
 try{await runtime.ready;manager=new ImageManagement(runtime,db,home);await manager.ready
  const manifest={protocol:1 as const,imageId:next,architecture:'arm64' as const,driver:'0.20.0' as const,layer:'5' as const,createdAt:1}
  vi.spyOn(manager.images,'inspect').mockResolvedValue(manifest)
  vi.spyOn(runtime.provider,'verifyRuntime').mockResolvedValue();vi.spyOn(runtime.provider,'ensure').mockResolvedValue({id:'fixture',containerId:'fixture',workspace:'/fixture',running:true,isolation:'container'});vi.spyOn(runtime.provider,'health').mockResolvedValue({});vi.spyOn(runtime.provider,'capture').mockResolvedValue({data:'fixture',width:1,height:1,capturedAt:1});vi.spyOn(runtime.provider,'remove').mockResolvedValue();vi.spyOn(runtime.provider,'deleteWorkspace').mockResolvedValue()
  const release=runtime.pool.beginMaintenance();await expect(manager.start('owner',randomUUID(),'activate',{imageId:next},async()=>{})).rejects.toMatchObject({code:'computer_busy'});release()
  const id=randomUUID();await manager.start('owner',id,'activate',{imageId:next},async()=>{});await vi.waitFor(()=>expect(runtime.config.imageId).toBe(next));await manager.close()
  expect(JSON.parse(String(db.prepare('SELECT value FROM managed_image_selection').get()!.value))).toMatchObject({active:next,previous:initial,base:initial})
  config.computers.imageId=initial
  const resumed=new ImageManagement(runtime,db,home);await resumed.ready;expect(runtime.config.imageId).toBe(next);await resumed.close()
 }finally{await manager?.close();await runtime.controls.close();await runtime.pool.close();db.close();await rm(home,{recursive:true,force:true})}
})
it('uploads ordered bounded chunks, detects conflicting retries and isolates the uploader',async()=>{
 const home=await mkdtemp(join(tmpdir(),'yaoyao-images-upload-')),db=new DatabaseSync(':memory:'),runtime=new ComputerRuntime(db,{protocol:1,runnerId:randomUUID(),serverURL:'http://127.0.0.1',hermesURL:'http://127.0.0.1',token:'fixture',allowedProfiles:['default'],artifactRoots:[],computers:{runtime:'docker',imageId:'sha256:'+'a'.repeat(64),python:'/fixture',hermesHome:'/fixture',hermesSource:'/fixture'}},home)
 const manager=new ImageManagement(runtime,db,home)
 try{await manager.ready;const id=randomUUID(),data=Buffer.from('image-fixture');const path=join(home,'hash');await writeFile(path,data)
  await manager.beginUpload('owner',id,data.length,{protocol:1,imageId:'sha256:'+'a'.repeat(64),architecture:'arm64',driver:'0.20.0',layer:'5',createdAt:1,archiveSha256:await fileSHA256(path)})
  const bytes=data.toString('base64');expect(await manager.appendUpload('owner',id,0,bytes)).toEqual({offset:data.length});expect(await manager.appendUpload('owner',id,0,bytes)).toEqual({offset:data.length})
  await expect(manager.appendUpload('other',id,0,bytes)).rejects.toMatchObject({code:'image_upload_missing'})
  await expect(manager.appendUpload('owner',id,0,Buffer.from('wrong').toString('base64'))).rejects.toThrow()
  await manager.cancelUpload('owner',id)
 }finally{await manager.close();await runtime.controls.close();await runtime.pool.close();db.close();await rm(home,{recursive:true,force:true})}
})
it('keeps the node in maintenance until uncertain validation cleanup is confirmed',async()=>{
 const home=await mkdtemp(join(tmpdir(),'yaoyao-images-cleanup-')),db=new DatabaseSync(':memory:'),initial='sha256:'+'a'.repeat(64)
 const runtime=new ComputerRuntime(db,{protocol:1,runnerId:randomUUID(),serverURL:'http://127.0.0.1',hermesURL:'http://127.0.0.1',token:'fixture',allowedProfiles:['default'],artifactRoots:[],computers:{runtime:'docker',imageId:initial,python:'/fixture',hermesHome:'/fixture',hermesSource:'/fixture'}},home),manager=new ImageManagement(runtime,db,home)
 let stopped=false
 try{await manager.ready;vi.spyOn(manager.images,'inspect').mockResolvedValue({protocol:1,imageId:'sha256:'+'b'.repeat(64),architecture:'arm64',driver:'0.20.0',layer:'5',createdAt:1});vi.spyOn(manager.images,'run').mockResolvedValue('');vi.spyOn(runtime.provider,'verifyRuntime').mockResolvedValue();vi.spyOn(runtime.provider,'ensure').mockResolvedValue({id:'fixture',containerId:'fixture',running:true,workspace:'/fixture',isolation:'container'});vi.spyOn(runtime.provider,'health').mockResolvedValue({});vi.spyOn(runtime.provider,'capture').mockResolvedValue({data:'fixture',width:1,height:1,capturedAt:1});vi.spyOn(runtime.provider,'remove').mockImplementation(async()=>{if(!stopped)throw new Error('stop uncertain')});vi.spyOn(runtime.provider,'deleteWorkspace').mockResolvedValue()
  await manager.start('owner',randomUUID(),'activate',{imageId:'sha256:'+'b'.repeat(64)},async()=>{})
  await vi.waitFor(async()=>expect((await manager.inventory('owner')).jobs[0]?.state).toBe('uncertain'))
  expect(()=>runtime.pool.beginMaintenance()).toThrow('请先停止');expect(runtime.config.imageId).toBe(initial)
  stopped=true;expect((await manager.inventory('owner')).busy).toBe(false)
  const release=runtime.pool.beginMaintenance();release()
 }finally{stopped=true;await manager.close();await runtime.controls.close();await runtime.pool.close();db.close();await rm(home,{recursive:true,force:true})}
})