ImpTune/tests/test_packages.py

"""Integration tests for package export endpoints — NinjaRMM ZIP and .intunewin."""
import io
import json
import zipfile

import pytest


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


@pytest.fixture
def driver_zip_bytes():
    """Create a minimal valid driver ZIP with a fake INF and CAT file."""
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w", compression=zipfile.ZIP_DEFLATED) as zf:
        zf.writestr("printer.inf", "[Version]\nSignature=$WINDOWS NT$\n")
        zf.writestr("printer.cat", "FAKE_CAT")
    return buf.getvalue()


@pytest.fixture
def setup_printer_with_driver(tmp_data_dir, driver_zip_bytes):
    """Create Driver record (with ZIP on disk) and Printer record linked to it."""
    import hashlib
    import os

    import imptune.config as cfg
    from imptune.db.models import Driver, Printer

    # Write driver ZIP to DRIVERS_DIR
    sha = hashlib.sha256(driver_zip_bytes).hexdigest()
    drivers_dir = cfg.DRIVERS_DIR
    os.makedirs(drivers_dir, exist_ok=True)
    zip_path = os.path.join(drivers_dir, f"{sha}.zip")
    with open(zip_path, "wb") as f:
        f.write(driver_zip_bytes)

    driver = Driver.create(
        sha256=sha,
        original_filename="printer_driver.zip",
        size_bytes=len(driver_zip_bytes),
        driver_desc=json.dumps(["HP LaserJet Pro"]),
        inf_filename="printer.inf",
    )
    printer = Printer.create(
        name="Test Printer",
        ip_address="192.168.1.100",
        port_name="IP_192.168.1.100",
        driver=driver,
        duplex_mode="OneSided",
        color_mode=True,
        paper_size="A4",
        collate=True,
    )
    return printer, driver


@pytest.fixture
def printer_no_driver(tmp_data_dir):
    """Create Printer record with no driver assigned."""
    from imptune.db.models import Printer

    return Printer.create(
        name="No Driver Printer",
        ip_address="10.0.0.1",
        port_name="IP_10.0.0.1",
        driver=None,
    )


# ---------------------------------------------------------------------------
# NinjaRMM ZIP endpoint
# ---------------------------------------------------------------------------


class TestNinjaDownload:
    def test_returns_zip(self, client, setup_printer_with_driver):
        printer, _ = setup_printer_with_driver
        resp = client.get(f"/printers/{printer.id}/packages/ninja")
        assert resp.status_code == 200
        assert resp.headers["content-type"] == "application/zip"
        assert "attachment" in resp.headers["content-disposition"]
        assert ".zip" in resp.headers["content-disposition"]

    def test_zip_contains_install_script(self, client, setup_printer_with_driver):
        printer, _ = setup_printer_with_driver
        resp = client.get(f"/printers/{printer.id}/packages/ninja")
        assert resp.status_code == 200
        safe_name = printer.name.replace(" ", "_")
        with zipfile.ZipFile(io.BytesIO(resp.content)) as zf:
            names = zf.namelist()
        assert f"{safe_name}/install.ps1" in names

    def test_zip_contains_driver_files(self, client, setup_printer_with_driver):
        printer, _ = setup_printer_with_driver
        resp = client.get(f"/printers/{printer.id}/packages/ninja")
        assert resp.status_code == 200
        safe_name = printer.name.replace(" ", "_")
        with zipfile.ZipFile(io.BytesIO(resp.content)) as zf:
            names = zf.namelist()
        # Driver ZIP contained printer.inf and printer.cat
        assert f"{safe_name}/drivers/printer.inf" in names
        assert f"{safe_name}/drivers/printer.cat" in names

    def test_404_missing_printer(self, client, tmp_data_dir):
        resp = client.get("/printers/9999/packages/ninja")
        assert resp.status_code == 404

    def test_422_no_driver(self, client, printer_no_driver):
        resp = client.get(f"/printers/{printer_no_driver.id}/packages/ninja")
        assert resp.status_code == 422


# ---------------------------------------------------------------------------
# .intunewin endpoint
# ---------------------------------------------------------------------------


class TestIntunewinDownload:
    def test_returns_intunewin(self, client, setup_printer_with_driver):
        printer, _ = setup_printer_with_driver
        resp = client.get(f"/printers/{printer.id}/packages/intunewin")
        assert resp.status_code == 200
        assert resp.headers["content-type"] == "application/octet-stream"
        assert "attachment" in resp.headers["content-disposition"]
        assert ".intunewin" in resp.headers["content-disposition"]

    def test_intunewin_is_valid_zip(self, client, setup_printer_with_driver):
        """Outer .intunewin file must be a valid ZIP with IntuneWinPackage/ structure."""
        printer, _ = setup_printer_with_driver
        resp = client.get(f"/printers/{printer.id}/packages/intunewin")
        assert resp.status_code == 200
        with zipfile.ZipFile(io.BytesIO(resp.content)) as zf:
            names = zf.namelist()
        assert any(n.startswith("IntuneWinPackage/") for n in names)

    def test_404_missing_printer(self, client, tmp_data_dir):
        resp = client.get("/printers/9999/packages/intunewin")
        assert resp.status_code == 404

    def test_422_no_driver(self, client, printer_no_driver):
        resp = client.get(f"/printers/{printer_no_driver.id}/packages/intunewin")
        assert resp.status_code == 422