from pathlib import Path
import tempfile
import unittest
from unittest.mock import patch
from subprocess import CompletedProcess

from restore_drill import DatabaseDrill
from transport import OperationsError


class DatabaseDrillTest(unittest.TestCase):
    def setUp(self):
        parent = Path(__file__).resolve().parent.parent / "storage/app/private/backup-check"
        parent.mkdir(parents=True, exist_ok=True)
        self.temporary = tempfile.TemporaryDirectory(prefix="drill-unit-", dir=parent)
        self.addCleanup(self.temporary.cleanup)
        self.root = Path(self.temporary.name)
        self.sql = self.root / "synthetic.sql"
        self.sql.write_text("CREATE TABLE example (id INT);\n")

    def test_drill_cannot_run_on_the_small_live_host(self):
        drill = DatabaseDrill(self.root)
        with patch.object(drill, "command", return_value=CompletedProcess([], 0, str(4 * 1024**3).encode())) as call:
            with self.assertRaises(OperationsError):
                drill.run(self.sql, "post76")
        self.assertEqual(1, call.call_count)

    def test_import_failure_removes_only_the_generated_recovery_container(self):
        drill = DatabaseDrill(self.root)
        calls = []

        def command(arguments, input_file=None, allowed_failure=False):
            calls.append(arguments)
            if arguments[0] == "info":
                return CompletedProcess([], 0, str(8 * 1024**3).encode())
            if input_file is not None:
                raise OperationsError("synthetic import failure")
            if "recovery-ready" in arguments:
                return CompletedProcess([], 0, b"1\n")
            return CompletedProcess([], 0, b"")

        with patch.object(drill, "command", side_effect=command):
            with self.assertRaises(OperationsError):
                drill.run(self.sql, "post76")
        create = next(arguments for arguments in calls if arguments[0] == "run")
        self.assertIn("none", create)
        self.assertNotIn("--publish", create)
        self.assertRegex(drill.name, r"^post76-restore-[a-f0-9]{32}$")
        self.assertEqual(["rm", "--force", "--volumes", drill.name], calls[-1])

    def test_untrusted_images_and_identifiers_are_rejected_before_docker(self):
        with self.assertRaises(OperationsError):
            DatabaseDrill(self.root, "untrusted/image:latest")
        drill = DatabaseDrill(self.root)
        with patch.object(drill, "command") as command:
            with self.assertRaises(OperationsError):
                drill.run(self.sql, "post76; arbitrary sql")
        command.assert_not_called()

    def test_import_waits_for_authenticated_database_readiness(self):
        drill = DatabaseDrill(self.root)
        ready_attempts = []

        def command(arguments, input_file=None, allowed_failure=False):
            if arguments[0] == "info":
                return CompletedProcess([], 0, str(8 * 1024**3).encode())
            if "recovery-ready" in arguments:
                ready_attempts.append(arguments)
                return CompletedProcess([], 0, b"1\n" if len(ready_attempts) == 2 else b"")
            if input_file is not None:
                self.assertEqual(2, len(ready_attempts))
                self.assertEqual(self.sql.read_bytes(), input_file.read())
            if "recovery-check" in arguments:
                return CompletedProcess([], 0, b"example\n")
            return CompletedProcess([], 0, b"")

        with patch.object(drill, "command", side_effect=command), patch("restore_drill.time.sleep") as sleep:
            result = drill.run(self.sql, "post76")
        self.assertTrue(result["database_imported"])
        self.assertEqual(1, result["tables"])
        sleep.assert_called_once_with(2)


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