OctoPrint-BambuPrinter/test/test_gcode_execution.py

from __future__ import annotations
from collections.abc import Callable
from datetime import datetime, timezone
import logging
from pathlib import Path
import time
from typing import Any
import unittest
from unittest.mock import MagicMock
import unittest.mock

from octoprint_bambu_printer.bambu_print_plugin import BambuPrintPlugin
from octoprint_bambu_printer.printer.bambu_virtual_printer import BambuVirtualPrinter
from octoprint_bambu_printer.printer.remote_sd_card_file_list import (
    FileInfo,
    RemoteSDCardFileList,
)
from octoprint_bambu_printer.printer.states.idle_state import IdleState
from octoprint_bambu_printer.printer.states.paused_state import PausedState
from octoprint_bambu_printer.printer.states.print_finished_state import (
    PrintFinishedState,
)
from octoprint_bambu_printer.printer.states.printing_state import PrintingState
from pytest import fixture


@fixture
def output_test_folder(output_folder: Path):
    folder = output_folder / "test_gcode"
    folder.mkdir(parents=True, exist_ok=True)
    return folder


@fixture
def log_test():
    return logging.getLogger("gcode_unittest")


class DictGetter:
    def __init__(self, options: dict) -> None:
        self._options: dict[str | tuple[str, ...], Any] = options

    def __call__(self, key: str | list[str] | tuple[str, ...]):
        if isinstance(key, list):
            key = tuple(key)
        return self._options.get(key, None)


@fixture
def settings(output_test_folder):
    _settings = MagicMock()
    _settings.get.side_effect = DictGetter(
        {
            "serial": "BAMBU",
            "host": "localhost",
            "access_code": "12345",
        }
    )
    _settings.get_boolean.side_effect = DictGetter({"forceChecksum": False})

    log_file_path = output_test_folder / "log.txt"
    log_file_path.touch()
    _settings.get_plugin_logfile_path.return_value = log_file_path.as_posix()
    return _settings


@fixture
def profile_manager():
    _profile_manager = MagicMock()
    _profile_manager.get_current.side_effect = MagicMock()
    _profile_manager.get_current().get.side_effect = DictGetter(
        {
            "heatedChamber": False,
        }
    )
    return _profile_manager


@fixture
def files_info_ftp():
    def _f_date(dt: datetime):
        return dt.replace(tzinfo=timezone.utc).strftime("%Y%m%d%H%M%S")

    return {
        "print.3mf": (1000, _f_date(datetime(2024, 5, 6))),
        "print2.3mf": (1200, _f_date(datetime(2024, 5, 7))),
        "cache/print.3mf": (1200, _f_date(datetime(2024, 5, 7))),
        "cache/print2.3mf": (1200, _f_date(datetime(2024, 5, 7))),
    }


@fixture
def ftps_session_mock(files_info_ftp):
    with unittest.mock.patch(
        "octoprint_bambu_printer.printer.ftpsclient.ftpsclient.IoTFTPSClient"
    ) as ftps_client_mock:
        ftps_session = MagicMock()
        ftps_session.size.side_effect = DictGetter(
            {file: info[0] for file, info in files_info_ftp.items()}
        )

        ftps_session.sendcmd.side_effect = DictGetter(
            {f"MDTM {file}": info[1] for file, info in files_info_ftp.items()}
        )

        all_files = list(files_info_ftp.keys())
        file_registry = DictGetter(
            {
                ("", ".3mf"): list(
                    filter(lambda f: Path(f).parent == Path("."), all_files)
                ),
                ("cache/", ".3mf"): list(
                    filter(lambda f: Path(f).parent == Path("cache/"), all_files)
                ),
            }
        )
        ftps_client_mock.list_files.side_effect = lambda folder, ext: file_registry(
            (folder, ext)
        )
        ftps_client_mock.ftps_session = ftps_session
        RemoteSDCardFileList._connect_ftps_server = MagicMock(
            return_value=ftps_client_mock
        )
        yield


@fixture(scope="function")
def printer(output_test_folder, settings, profile_manager, log_test, ftps_session_mock):
    async def _mock_connection(self):
        pass

    BambuVirtualPrinter._create_connection_async = _mock_connection
    serial_obj = BambuVirtualPrinter(
        settings,
        profile_manager,
        data_folder=output_test_folder,
        serial_log_handler=log_test,
        read_timeout=0.01,
        faked_baudrate=115200,
    )
    serial_obj._bambu_client = MagicMock()
    yield serial_obj
    serial_obj.close()


def test_initial_state(printer: BambuVirtualPrinter):
    assert isinstance(printer.current_state, IdleState)


def test_list_sd_card(printer: BambuVirtualPrinter):
    printer.write(b"M20\n")  # GCode for listing SD card
    printer.flush()
    result = printer.readlines()
    assert result[0] == b"Begin file list"
    assert result[1].endswith(b'"print.3mf"')
    assert result[2].endswith(b'"print2.3mf"')
    assert result[3] == b"End file list"
    assert result[4] == b"ok"


def test_cannot_start_print_without_file(printer: BambuVirtualPrinter):
    printer.write(b"M24\n")
    printer.flush()
    result = printer.readlines()
    assert result[0] == b"ok"
    assert isinstance(printer.current_state, IdleState)


def test_non_existing_file_not_ok(printer: BambuVirtualPrinter):
    printer.write(b"M23 non_existing.3mf\n")
    printer.flush()
    result = printer.readlines()
    assert result[0] == b"ok"
    assert printer.file_system.selected_file is None


def test_print_started_with_selected_file(printer: BambuVirtualPrinter):
    assert printer.file_system.selected_file is None

    printer.write(b"M20\n")
    printer.flush()
    printer.readlines()

    printer.write(b"M23 print.3mf\n")
    printer.flush()
    result = printer.readlines()
    assert result[0] == b"ok"

    assert printer.file_system.selected_file is not None
    assert printer.file_system.selected_file.file_name == "print.3mf"

    printer.write(b"M24\n")
    printer.flush()

    result = printer.readlines()
    assert result[0] == b"ok"
    assert isinstance(printer.current_state, PrintingState)


def test_pause_print(printer: BambuVirtualPrinter):
    gcode = b"G28\nG1 X10 Y10\n"
    printer.write(gcode)
    printer.write(b"M25\n")  # GCode for pausing the print
    result = printer.readline()
    assert isinstance(printer.current_state, PausedState)


def test_get_printing_info(printer: BambuVirtualPrinter):
    gcode = b"G28\nG1 X10 Y10\n"
    printer.write(gcode)
    printer.write(b"M27\n")  # GCode for getting printing info
    result = printer.readline()
    assert result == ""


def test_abort_print(printer: BambuVirtualPrinter):
    gcode = b"G28\nG1 X10 Y10\n"
    printer.write(gcode)
    printer.write(b"M26\n")  # GCode for aborting the print
    result = printer.readline()
    assert isinstance(printer.current_state, IdleState)


def test_print_finished(printer: BambuVirtualPrinter):
    gcode = b"G28\nG1 X10 Y10\n"
    printer.write(gcode)
    result = printer.readline()
    assert isinstance(printer.current_state, PrintFinishedState)