Home-Assistant/esphome/bambu-display/test_generate.py

"""Run with: python -m unittest discover -s esphome/bambu-display -v."""

import copy
import unittest

import yaml

from generate import HERE, generate, validate_printers


class FleetTests(unittest.TestCase):
    def setUp(self):
        self.printers = yaml.safe_load((HERE / 'printers.yaml').read_text())['printers']

    def test_all_subscriptions_match_repository_entities(self):
        known = {line.split(' | ')[0] for line in
                 (HERE.parent / 'ha-entities.txt').read_text().splitlines()}
        package = generate(self.printers)
        subscribed = {entry['entity_id'] for domain in ('sensor', 'text_sensor')
                      for entry in package[domain]}
        self.assertFalse(subscribed - known)

    def test_fleet_can_expand_beyond_configured_printers(self):
        expected_count = len(self.printers) + 2
        for index in (8, 9):
            extra = copy.deepcopy(self.printers[0])
            extra.update(id=f'printer_{index}', entity_prefix=f'new_printer_{index}')
            self.printers.append(extra)
        package = generate(self.printers)
        self.assertEqual(len([page for page in package['lvgl']['pages'] if not page.get('skip')]), expected_count)
        self.assertEqual(len(package['script']), expected_count)
        subscriptions = {entry['id']: entry['entity_id'] for entry in package['sensor']}
        self.assertEqual(subscriptions['printer_9_progress'], 'sensor.new_printer_9_print_progress')
        for page in package['lvgl']['pages']:
            self.assertIn('lvgl.page.next', page['on_swipe_left'][-1])
            self.assertIn('lvgl.page.previous', page['on_swipe_right'][-1])

    def test_elegoo_profiles_use_available_metrics_and_domains(self):
        package = generate(self.printers)
        subscriptions = {entry['id']: entry for domain in ('sensor', 'text_sensor')
                         for entry in package[domain]}
        self.assertEqual(subscriptions['centauri_carbon_nozzle_target']['entity_id'],
                         'number.centauri_carbon_target_nozzle_temp')
        self.assertEqual(subscriptions['centauri_carbon_fan']['entity_id'],
                         'sensor.centauri_carbon_model_fan_speed')
        for pid in ('centauri_carbon', 'saturn'):
            self.assertEqual(subscriptions[f'{pid}_online']['entity_id'], f'binary_sensor.{pid}_sdcp_status')
            self.assertEqual(subscriptions[f'{pid}_remaining_unit']['entity_id'], f'sensor.{pid}_remaining_print_time')
            self.assertEqual(subscriptions[f'{pid}_print_error']['entity_id'], f'sensor.{pid}_print_error')
            self.assertEqual(subscriptions[f'{pid}_print_error_reason']['entity_id'], f'sensor.{pid}_error_status_reason')
            self.assertNotIn(f'{pid}_hms_error', subscriptions)
        self.assertNotIn('saturn_nozzle', subscriptions)
        self.assertNotIn('saturn_fan', subscriptions)
        self.assertEqual(subscriptions['saturn_uv_unit']['entity_id'], 'sensor.saturn_uv_led_temp')
        self.assertEqual(subscriptions['saturn_elapsed_unit']['entity_id'], 'sensor.saturn_current_print_time')

    def test_resin_page_contains_only_resin_gauges(self):
        package = generate(self.printers)
        page = next(page for page in package['lvgl']['pages'] if page['id'] == 'saturn_page')
        arcs = {widget['arc']['id'] for widget in page['widgets'] if 'arc' in widget}
        self.assertEqual(arcs, {'saturn_progress_arc', 'saturn_uv_arc', 'saturn_layer_arc'})

    def test_number_overrides_are_supported(self):
        self.printers[-2]['entities'] = {'nozzle_target': 'number.custom_target'}
        package = generate(self.printers)
        entry = next(row for row in package['sensor'] if row['id'] == 'centauri_carbon_nozzle_target')
        self.assertEqual(entry['entity_id'], 'number.custom_target')

    def test_reject_unknown_profiles(self):
        self.printers[0]['profile'] = 'unsupported'
        with self.assertRaisesRegex(ValueError, 'unsupported profile'):
            generate(self.printers)

    def test_overrides_apply_to_values_and_unit_subscriptions(self):
        self.printers[0]['entities'] = {'time_left': 'sensor.custom_remaining'}
        package = generate(self.printers)
        for domain, key in [('sensor', 'time_left'), ('text_sensor', 'remaining_unit')]:
            entry = next(row for row in package[domain] if row['id'] == f'prints_charming_{key}')
            self.assertEqual(entry['entity_id'], 'sensor.custom_remaining')

    def test_orion_has_both_ams_units_with_their_actual_entity_prefixes(self):
        package = generate(self.printers)
        types = [entry['entity_id'] for entry in package['text_sensor']
                 if entry['id'].startswith('constipation_orion_tray_') and entry.get('attribute') == 'type']
        self.assertEqual(len(types), 8)
        self.assertIn('sensor.constipation_orion_ams_1_tray_1', types)
        self.assertIn('sensor.h2c_31b8ap5a2200288_ams_2_tray_4', types)

    def test_health_messages_import_numbered_ha_attributes(self):
        package = generate(self.printers)
        subscriptions = {entry['id']: entry for entry in package['text_sensor']}
        message = subscriptions['prints_charming_hms_1_message']
        self.assertEqual(message['entity_id'], 'binary_sensor.p1s_01p09c470102673_hms_errors')
        self.assertEqual(message['attribute'], '1-Error')
        self.assertEqual(subscriptions['prints_charming_hms_3_severity']['attribute'], '3-Severity')

    def test_reject_empty_duplicate_and_malformed_printer_ids(self):
        for invalid in ([], self.printers + [self.printers[0]],
                        [{**self.printers[0], 'id': 'invalid-id'}]):
            with self.subTest(invalid=invalid), self.assertRaises(ValueError):
                validate_printers(invalid)

    def test_reject_overflowing_filament_strip(self):
        self.printers[0]['external_spools'] = [{'label': 'EXT', 'entity': 'sensor.external_spool'}]
        with self.assertRaisesRegex(ValueError, 'eight slots'):
            generate(self.printers)

    def test_progress_and_dual_nozzle_widgets_fit_the_panel(self):
        package = generate(self.printers)
        for page in package['lvgl']['pages']:
            for widget in page['widgets']:
                props = next(iter(widget.values()))
                self.assertGreaterEqual(props.get('x', 0), 0)
                self.assertGreaterEqual(props.get('y', 0), 0)
                self.assertLessEqual(props.get('x', 0) + props.get('width', 0), 800)
                self.assertLessEqual(props.get('y', 0) + props.get('height', 0), 480)

    def test_rgb565_primaries_reach_the_correct_panel_bits(self):
        # mipi_rgb emits named RGB data pins with bytes exchanged, then copies
        # LVGL RGB565 bytes straight into the ESP32 little-endian framebuffer.
        # Exercise red, green, blue and combined primaries through that path.
        class Loader(yaml.SafeLoader):
            pass
        for tag in ('!secret', '!include'):
            Loader.add_constructor(tag, lambda loader, node: loader.construct_scalar(node))
        config = yaml.load((HERE.parent / 'bambu-printer-display.yaml').read_text(), Loader=Loader)
        display = config['display'][0]
        pins = display['data_pins']
        natural = pins['blue'] + pins['green'] + pins['red']
        driver = natural[8:] + natural[:8]
        order = {'big_endian': 'big', 'little_endian': 'little'}[config['lvgl']['byte_order']]
        for color in (0xF800, 0x07E0, 0x001F, 0xFFFF, 0x07FF, 0xF81F):
            memory_word = int.from_bytes(color.to_bytes(2, order), 'little')
            panel_word = sum(((memory_word >> i) & 1) << natural.index(pin)
                             for i, pin in enumerate(driver))
            self.assertEqual(panel_word, color, f'RGB565 color {color:#06x} is wired incorrectly')

    def test_diagnostic_color_page_is_skipped_during_normal_navigation(self):
        package = generate(self.printers)
        page = next(page for page in package['lvgl']['pages'] if page['id'] == 'color_test_page')
        self.assertTrue(page['skip'])


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