Skip to content

Commit 681acee

Browse files
shaggysashaggysa
committed
add a few unit tests related to stay_connected_menu
1 parent d7cd4bd commit 681acee

File tree

1 file changed

+313
-0
lines changed

1 file changed

+313
-0
lines changed

‎tests/test_cli.py‎

Lines changed: 313 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,22 @@
11
"""Tests for the pybricksdev CLI commands."""
22

33
import argparse
4+
import asyncio
45
import contextlib
56
import io
67
import os
8+
import subprocess
79
import tempfile
810
from unittest.mock import AsyncMock, Mock, mock_open, patch
911

1012
import pytest
13+
from packaging.version import Version
1114

1215
from pybricksdev.cli import Compile, Run, Tool, Udev
16+
from pybricksdev.connections.pybricks import (
17+
HubDisconnectError,
18+
HubPowerButtonPressedError,
19+
)
1320

1421

1522
class TestTool:
@@ -442,6 +449,312 @@ async def test_run_connection_error(self):
442449
# Verify disconnect was not called since connection failed
443450
mock_hub.disconnect.assert_not_called()
444451

452+
@pytest.mark.asyncio
453+
async def test_run_syntax_error(self):
454+
"""Test that the stay connected menu is called upon a syntax error when the appropriate flag is active."""
455+
456+
# Create a mock hub
457+
mock_hub = AsyncMock()
458+
mock_hub.run = AsyncMock(
459+
side_effect=subprocess.CalledProcessError(
460+
returncode=1, cmd="test", stderr=b"test"
461+
)
462+
)
463+
mock_hub.connect = AsyncMock()
464+
465+
# Set up mocks using ExitStack
466+
with contextlib.ExitStack() as stack:
467+
# Create and manage temporary file
468+
temp = stack.enter_context(
469+
tempfile.NamedTemporaryFile(
470+
suffix=".py", mode="w+", delete=False, encoding="utf-8"
471+
)
472+
)
473+
temp.write("print('test')")
474+
temp_path = temp.name
475+
stack.callback(os.unlink, temp_path)
476+
477+
# Create args
478+
args = argparse.Namespace(
479+
conntype="ble",
480+
file=stack.enter_context(open(temp_path, "r", encoding="utf-8")),
481+
name="MyHub",
482+
start=True,
483+
wait=True,
484+
stay_connected=True,
485+
)
486+
487+
mock_hub_class = stack.enter_context(
488+
patch(
489+
"pybricksdev.connections.pybricks.PybricksHubBLE",
490+
return_value=mock_hub,
491+
)
492+
)
493+
stack.enter_context(
494+
patch("pybricksdev.ble.find_device", return_value="mock_device")
495+
)
496+
mock_menu = stack.enter_context(
497+
patch("pybricksdev.cli.Run.stay_connected_menu")
498+
)
499+
500+
# Run the command
501+
run_cmd = Run()
502+
await run_cmd.run(args)
503+
504+
# Verify the hub was created and used correctly
505+
mock_hub_class.assert_called_once_with("mock_device")
506+
mock_hub.connect.assert_called_once()
507+
mock_hub.run.assert_called_once_with(temp_path, True)
508+
mock_menu.assert_called_once_with(mock_hub, args)
509+
mock_hub.disconnect.assert_called_once()
510+
511+
@pytest.mark.asyncio
512+
async def test_stay_connected_menu_integration(self):
513+
"""Test all of the basic options in the stay_connected menu."""
514+
515+
async def passthrough_awaitable(awaitable):
516+
return await awaitable
517+
518+
# Create a mock hub
519+
mock_hub = AsyncMock()
520+
mock_hub.fw_version = Version("3.2.0-beta.4")
521+
mock_hub.run = AsyncMock()
522+
mock_hub.connect = AsyncMock()
523+
mock_hub.start_user_program = AsyncMock()
524+
mock_hub._wait_for_user_program_stop = AsyncMock()
525+
mock_hub.race_disconnect = mock_hub.race_power_button_press = AsyncMock(
526+
side_effect=passthrough_awaitable
527+
)
528+
mock_hub.download = AsyncMock()
529+
530+
# create a mock questionary menu
531+
mock_selector = AsyncMock()
532+
mock_selector.ask_async.side_effect = [
533+
"Recompile and Run",
534+
"Recompile and Download",
535+
"Run Stored Program",
536+
"Exit",
537+
]
538+
539+
# Set up mocks using ExitStack
540+
with contextlib.ExitStack() as stack:
541+
# Create and manage temporary file
542+
temp = stack.enter_context(
543+
tempfile.NamedTemporaryFile(
544+
suffix=".py", mode="w+", delete=False, encoding="utf-8"
545+
)
546+
)
547+
temp.write("print('test')")
548+
temp_path = temp.name
549+
stack.callback(os.unlink, temp_path)
550+
551+
# Create args
552+
args = argparse.Namespace(
553+
conntype="ble",
554+
file=stack.enter_context(open(temp_path, "r", encoding="utf-8")),
555+
name="MyHub",
556+
start=True,
557+
wait=True,
558+
stay_connected=True,
559+
)
560+
561+
mock_hub_class = stack.enter_context(
562+
patch(
563+
"pybricksdev.connections.pybricks.PybricksHubBLE",
564+
return_value=mock_hub,
565+
)
566+
)
567+
stack.enter_context(
568+
patch("pybricksdev.ble.find_device", return_value="mock_device")
569+
)
570+
mock_selector = stack.enter_context(
571+
patch("questionary.select", return_value=mock_selector)
572+
)
573+
574+
# Run the command
575+
run_cmd = Run()
576+
await run_cmd.run(args)
577+
578+
# Verify the hub was created and used correctly
579+
mock_hub_class.assert_called_once_with("mock_device")
580+
mock_hub.connect.assert_called_once()
581+
assert mock_hub.run.call_count == 2
582+
mock_hub.run.assert_called_with(temp_path, wait=True)
583+
mock_hub.download.assert_called_once_with(temp_path)
584+
mock_hub.start_user_program.assert_called_once()
585+
mock_hub._wait_for_user_program_stop.assert_called_once()
586+
assert mock_selector.call_count == 4
587+
mock_hub.disconnect.assert_called_once()
588+
589+
@pytest.mark.asyncio
590+
async def test_stay_connected_menu_change_target_file(self):
591+
"""Test the change target file option."""
592+
593+
async def passthrough_awaitable(awaitable):
594+
return await awaitable
595+
596+
# Create a mock hub
597+
mock_hub = AsyncMock()
598+
mock_hub.run = AsyncMock()
599+
mock_hub.connect = AsyncMock()
600+
mock_hub.race_disconnect = mock_hub.race_power_button_press = AsyncMock(
601+
side_effect=passthrough_awaitable
602+
)
603+
mock_hub.download = AsyncMock()
604+
605+
# create a mock questionary menu
606+
mock_menu = AsyncMock()
607+
mock_menu.ask_async.side_effect = [
608+
"Change Target File",
609+
"Recompile and Run",
610+
"Exit",
611+
]
612+
613+
# Set up mocks using ExitStack
614+
with contextlib.ExitStack() as stack:
615+
# Create and manage temporary file
616+
temp = stack.enter_context(
617+
tempfile.NamedTemporaryFile(
618+
suffix=".py", mode="w+", delete=False, encoding="utf-8"
619+
)
620+
)
621+
new_target_file = stack.enter_context(
622+
tempfile.NamedTemporaryFile(
623+
suffix=".py", mode="w+", delete=False, encoding="utf-8"
624+
)
625+
)
626+
temp.write("print('test')")
627+
temp_path = temp.name
628+
stack.callback(os.unlink, temp_path)
629+
new_target_file.write("print('test')")
630+
new_path = new_target_file.name
631+
stack.callback(os.unlink, new_path)
632+
633+
# Create args
634+
args = argparse.Namespace(
635+
conntype="ble",
636+
file=stack.enter_context(open(temp_path, "r", encoding="utf-8")),
637+
name="MyHub",
638+
start=True,
639+
wait=True,
640+
stay_connected=True,
641+
)
642+
643+
stack.enter_context(
644+
patch("pybricksdev.ble.find_device", return_value="mock_device")
645+
)
646+
mock_selector = stack.enter_context(
647+
patch("questionary.select", return_value=mock_menu)
648+
)
649+
mock_file_selector = stack.enter_context(patch("questionary.path"))
650+
path_obj = AsyncMock()
651+
path_obj.ask_async = AsyncMock(return_value=new_path)
652+
mock_file_selector.return_value = path_obj
653+
654+
# Run the command
655+
run_cmd = Run()
656+
await run_cmd.stay_connected_menu(mock_hub, args)
657+
658+
assert mock_selector.call_count == 3
659+
mock_file_selector.assert_called_once()
660+
mock_hub.download.assert_called_once_with(new_path)
661+
mock_hub.run.assert_called_once_with(new_path, wait=True)
662+
663+
@pytest.mark.asyncio
664+
async def test_stay_connected_menu_interruptions(self):
665+
"""Test the stay_connected_menu being interrupted by a power button press or hub disconnect."""
666+
disconnect_call_count = 0
667+
power_call_count = 0
668+
669+
# simulates the hub disconnecting on the first call,
670+
async def mock_race_disconnect(awaitable):
671+
task = asyncio.ensure_future(awaitable)
672+
nonlocal disconnect_call_count
673+
disconnect_call_count += 1
674+
if disconnect_call_count == 1:
675+
task.cancel()
676+
raise HubDisconnectError("hub disconnected")
677+
return await awaitable
678+
679+
async def mock_race_power_button_press(awaitable):
680+
task = asyncio.ensure_future(awaitable)
681+
nonlocal power_call_count
682+
power_call_count += 1
683+
if power_call_count == 2:
684+
task.cancel()
685+
raise HubPowerButtonPressedError("power button pressed")
686+
return await awaitable
687+
688+
# Create a mock hub
689+
mock_hub = AsyncMock()
690+
mock_hub.run = AsyncMock()
691+
mock_hub.connect = AsyncMock()
692+
mock_hub.disconnect = AsyncMock()
693+
mock_hub.race_disconnect = AsyncMock(
694+
side_effect=mock_race_disconnect,
695+
)
696+
mock_hub.race_power_button_press = AsyncMock(
697+
side_effect=mock_race_power_button_press,
698+
)
699+
mock_hub._wait_for_power_button_release = AsyncMock()
700+
mock_hub._wait_for_user_program_stop = AsyncMock()
701+
# create a mock questionary menu
702+
mock_menu = AsyncMock()
703+
mock_menu.ask_async.side_effect = [
704+
"Recompile and Run",
705+
"Exit",
706+
]
707+
mock_confirm = AsyncMock()
708+
mock_confirm.ask_async.return_value = True
709+
710+
# Set up mocks using ExitStack
711+
with contextlib.ExitStack() as stack:
712+
# Create and manage temporary file
713+
temp = stack.enter_context(
714+
tempfile.NamedTemporaryFile(
715+
suffix=".py", mode="w+", delete=False, encoding="utf-8"
716+
)
717+
)
718+
temp.write("print('test')")
719+
temp_path = temp.name
720+
stack.callback(os.unlink, temp_path)
721+
722+
# Create args
723+
args = argparse.Namespace(
724+
conntype="ble",
725+
file=stack.enter_context(open(temp_path, "r", encoding="utf-8")),
726+
name="MyHub",
727+
start=True,
728+
wait=True,
729+
stay_connected=True,
730+
)
731+
732+
mock_hub_class = stack.enter_context(
733+
patch(
734+
"pybricksdev.connections.pybricks.PybricksHubBLE",
735+
return_value=mock_hub,
736+
)
737+
)
738+
739+
stack.enter_context(
740+
patch("pybricksdev.ble.find_device", return_value="mock_device")
741+
)
742+
mock_selector = stack.enter_context(
743+
patch("questionary.select", return_value=mock_menu)
744+
)
745+
stack.enter_context(patch("questionary.confirm", return_value=mock_confirm))
746+
747+
# Run the command
748+
run_cmd = Run()
749+
await run_cmd.stay_connected_menu(mock_hub, args)
750+
751+
assert mock_selector.call_count == 4
752+
mock_hub_class.assert_called_once()
753+
mock_hub.connect.assert_called_once()
754+
mock_hub._wait_for_power_button_release.assert_called_once()
755+
mock_hub._wait_for_user_program_stop.assert_called_once()
756+
mock_hub.run.assert_called_once_with(temp_path, wait=True)
757+
445758

446759
class TestCompile:
447760
"""Tests for the Compile command."""

0 commit comments

Comments
 (0)