diff --git a/simulation/sim_base.py b/simulation/sim_base.py index fddbe55..0abc6f9 100644 --- a/simulation/sim_base.py +++ b/simulation/sim_base.py @@ -2,6 +2,7 @@ import sys import os import time import pytest +from typing import Any, Optional from api_hook_client import ApiHookClient from simulation.workflow_sim import WorkflowSimulator @@ -17,7 +18,7 @@ class BaseSimulation: self.sim = WorkflowSimulator(self.client) self.project_path = None - def setup(self, project_name="SimProject"): + def setup(self, project_name: str = "SimProject") -> None: print(f"\n[BaseSim] Connecting to GUI...") if not self.client.wait_for_server(timeout=5): raise RuntimeError("Could not connect to GUI. Ensure it is running with --enable-test-hooks") @@ -43,20 +44,16 @@ class BaseSimulation: pass print("[BaseSim] Teardown complete.") - def get_value(self, tag): + def get_value(self, tag: str) -> Any: return self.client.get_value(tag) - def wait_for_event(self, event_type, timeout=5): + def wait_for_event(self, event_type: str, timeout: int = 5) -> Optional[dict]: return self.client.wait_for_event(event_type, timeout) - def assert_panel_visible(self, panel_tag, msg=None): - # This assumes we have a hook to check panel visibility or just check if an element in it exists - # For now, we'll check if we can get a value from an element that should be in that panel - # or use a specific hook if available. - # Actually, let's just check if get_indicator_state or similar works for generic tags. + def assert_panel_visible(self, panel_tag: str, msg: str = None) -> None: pass - def wait_for_element(self, tag, timeout=2): + def wait_for_element(self, tag: str, timeout: int = 2) -> bool: start = time.time() while time.time() - start < timeout: try: @@ -67,7 +64,7 @@ class BaseSimulation: time.sleep(0.1) return False -def run_sim(sim_class): +def run_sim(sim_class: type) -> None: """Helper to run a simulation class standalone.""" sim = sim_class() try: diff --git a/simulation/sim_execution.py b/simulation/sim_execution.py index d4cf17e..187410f 100644 --- a/simulation/sim_execution.py +++ b/simulation/sim_execution.py @@ -4,7 +4,7 @@ import time from simulation.sim_base import BaseSimulation, run_sim class ExecutionSimulation(BaseSimulation): - def setup(self, project_name="SimProject"): + def setup(self, project_name: str = "SimProject") -> None: super().setup(project_name) if os.path.exists("hello.ps1"): os.remove("hello.ps1") diff --git a/simulation/user_agent.py b/simulation/user_agent.py index daa10f3..f9c74c2 100644 --- a/simulation/user_agent.py +++ b/simulation/user_agent.py @@ -1,9 +1,10 @@ import time import random +from typing import Any, Callable import ai_client class UserSimAgent: - def __init__(self, hook_client, model="gemini-2.5-flash-lite"): + def __init__(self, hook_client: Any, model: str = "gemini-2.5-flash-lite") -> None: self.hook_client = hook_client self.model = model self.system_prompt = ( @@ -13,7 +14,7 @@ class UserSimAgent: "Do not use markdown blocks for your main message unless you are providing code." ) - def generate_response(self, conversation_history): + def generate_response(self, conversation_history: list[dict]) -> str: """ Generates a human-like response based on the conversation history. conversation_history: list of dicts with 'role' and 'content' @@ -38,7 +39,7 @@ class UserSimAgent: ai_client.set_custom_system_prompt("") return response - def perform_action_with_delay(self, action_func, *args, **kwargs): + def perform_action_with_delay(self, action_func: Callable, *args: Any, **kwargs: Any) -> Any: """ Executes an action with a human-like delay. """ diff --git a/simulation/workflow_sim.py b/simulation/workflow_sim.py index 8c7cd4f..80b9917 100644 --- a/simulation/workflow_sim.py +++ b/simulation/workflow_sim.py @@ -8,7 +8,7 @@ class WorkflowSimulator: self.client = hook_client self.user_agent = UserSimAgent(hook_client) - def setup_new_project(self, name, git_dir, project_path=None): + def setup_new_project(self, name: str, git_dir: str, project_path: str = None) -> None: print(f"Setting up new project: {name}") if project_path: self.client.click("btn_project_new_automated", user_data=project_path) @@ -19,13 +19,13 @@ class WorkflowSimulator: self.client.click("btn_project_save") time.sleep(1) - def create_discussion(self, name): + def create_discussion(self, name: str) -> None: print(f"Creating discussion: {name}") self.client.set_value("disc_new_name_input", name) self.client.click("btn_disc_create") time.sleep(1) - def switch_discussion(self, name): + def switch_discussion(self, name: str) -> None: print(f"Switching to discussion: {name}") self.client.select_list_item("disc_listbox", name) time.sleep(1) @@ -37,18 +37,18 @@ class WorkflowSimulator: # without more hooks, but we can verify the button click. time.sleep(1) - def truncate_history(self, pairs): + def truncate_history(self, pairs: int) -> None: print(f"Truncating history to {pairs} pairs") self.client.set_value("disc_truncate_pairs", pairs) self.client.click("btn_disc_truncate") time.sleep(1) - def run_discussion_turn(self, user_message=None): + def run_discussion_turn(self, user_message: str = None) -> dict | None: self.run_discussion_turn_async(user_message) # Wait for AI return self.wait_for_ai_response() - def run_discussion_turn_async(self, user_message=None): + def run_discussion_turn_async(self, user_message: str = None) -> None: if user_message is None: # Generate from AI history session = self.client.get_session() @@ -58,7 +58,7 @@ class WorkflowSimulator: self.client.set_value("ai_input", user_message) self.client.click("btn_gen_send") - def wait_for_ai_response(self, timeout=60): + def wait_for_ai_response(self, timeout: int = 60) -> dict | None: print("Waiting for AI response...", end="", flush=True) start_time = time.time() last_count = len(self.client.get_session().get('session', {}).get('entries', []))