#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-3.0-or-later
from __future__ import annotations
import argparse,struct,sys,tempfile,unittest
from pathlib import Path
P=argparse.ArgumentParser();P.add_argument("--root",type=Path,required=True);R=P.parse_args().root;sys.path.insert(0,str(R/"tools"))
from phase10dm_snapshot_protocol import ProtocolError  # noqa:E402
from phase10dn_snapshot_receiver import SnapshotReceiver  # noqa:E402
F=struct.Struct("<8sIIiiQQQQq");DATA=b"SQLite format 3\x00"+b"x"*16
def frame(k,ino=2):return F.pack(b"CHS10DM1",1,k,0,0,len(DATA),len(DATA) if k==2 else 0,1,ino,3)
class Tests(unittest.TestCase):
 def test_chunked_success(self):
  with tempfile.TemporaryDirectory() as d:
   p=Path(d)/"snapshot.db";r=SnapshotReceiver(p);raw=frame(1)+DATA+frame(2)
   for byte in raw:r.feed(bytes([byte]))
   self.assertEqual(r.finish().size,len(DATA));self.assertEqual(p.read_bytes(),DATA)
 def test_existing_rejected(self):
  with tempfile.TemporaryDirectory() as d:
   p=Path(d)/"x";p.write_bytes(b"x")
   with self.assertRaises(ProtocolError):SnapshotReceiver(p)
 def test_truncated_aborts(self):
  with tempfile.TemporaryDirectory() as d:
   r=SnapshotReceiver(Path(d)/"x");r.feed(frame(1)+DATA)
   with self.assertRaises(ProtocolError):r.finish()
   self.assertEqual(r.state,"ABORTED")
 def test_metadata_mismatch(self):
  with tempfile.TemporaryDirectory() as d:
   r=SnapshotReceiver(Path(d)/"x");r.feed(frame(1)+DATA+frame(2,ino=9))
   with self.assertRaises(ProtocolError):r.finish()
 def test_trailing_rejected(self):
  with tempfile.TemporaryDirectory() as d:
   r=SnapshotReceiver(Path(d)/"x")
   with self.assertRaises(ProtocolError):r.feed(frame(1)+DATA+frame(2)+b"x")
if __name__=="__main__":unittest.main(argv=[__file__])