From 94684db953c58ba5b3c819725cf60903725ee260 Mon Sep 17 00:00:00 2001 From: Richard Yu Date: Fri, 18 Jul 2025 17:14:46 -0700 Subject: [PATCH] fix lint --- execution.py | 34 +++++++++++++------------- tests/inference/test_execution.py | 40 +++++++++++++++---------------- 2 files changed, 37 insertions(+), 37 deletions(-) diff --git a/execution.py b/execution.py index 6804aa0e3..daf61ba19 100644 --- a/execution.py +++ b/execution.py @@ -1121,7 +1121,7 @@ class PromptQueue: Retrieves execution history in chronological order with pagination support. Returns a lightweight list of history objects. Used by the /history_v2. - + API Output Structure: { "history": [ @@ -1135,7 +1135,7 @@ class PromptQueue: "extra_data": dict # Additional metadata (workflow removed from extra_pnginfo) } | None, # None if no prompt data available "status": { - "status_str": str, # "success" | "error" + "status_str": str, # "success" | "error" "messages": [ # Filtered execution event messages (event_name: str, event_data: dict) ] @@ -1144,7 +1144,7 @@ class PromptQueue: # ... more history items ] } - + Parameters: - max_items: Maximum number of items to return (None = all) - offset: Starting index (0-based, negative values calculated from end) @@ -1161,36 +1161,36 @@ class PromptQueue: history_items = [] for key in selected_keys: history_entry = self.history[key] - + filtered_prompt = None if "prompt" in history_entry: priority, prompt_id, _, extra_data, _ = history_entry["prompt"] - + filtered_extra_data = {} for k, v in extra_data.items(): if k == "extra_pnginfo": filtered_extra_data[k] = { - pk: pv for pk, pv in v.items() + pk: pv for pk, pv in v.items() if pk != "workflow" } else: filtered_extra_data[k] = v - + filtered_prompt = { "priority": priority, "prompt_id": prompt_id, "extra_data": filtered_extra_data } - + status = None if history_entry.get("status"): status = { "status_str": history_entry["status"]["status_str"], - "messages": [(e, {k: v for k, v in d.items() if k != "nodes"}) + "messages": [(e, {k: v for k, v in d.items() if k != "nodes"}) if e == "execution_cached" else (e, d) for e, d in history_entry["status"]["messages"]] } - + item = { "prompt_id": key, "outputs": history_entry.get("outputs", {}), @@ -1198,7 +1198,7 @@ class PromptQueue: "prompt": filtered_prompt, "status": status } - + history_items.append(item) return {"history": history_items} @@ -1207,7 +1207,7 @@ class PromptQueue: """ Retrieves execution history for a specific prompt ID. Used by /history_v2/:prompt_id - + API Output Structure: { "": { @@ -1227,15 +1227,15 @@ class PromptQueue: } | None # None if no status recorded } } - + Returns empty dict {} if prompt_id not found. """ with self.mutex: if prompt_id in self.history: history_entry = self.history[prompt_id] - + new_entry = {} - + if "prompt" in history_entry: priority, prompt_id_inner, prompt_data, extra_data, outputs_to_execute = history_entry["prompt"] new_entry["prompt"] = { @@ -1245,11 +1245,11 @@ class PromptQueue: "extra_data": extra_data, "outputs_to_execute": outputs_to_execute } - + for key, value in history_entry.items(): if key != "prompt": new_entry[key] = value - + return {prompt_id: new_entry} else: return {} diff --git a/tests/inference/test_execution.py b/tests/inference/test_execution.py index 2184f544d..b5f41c70d 100644 --- a/tests/inference/test_execution.py +++ b/tests/inference/test_execution.py @@ -660,15 +660,15 @@ class TestExecution: # Test history_v2 endpoint for specific prompt specific_history = client.get_history_v2_for_prompt(prompt_id) assert prompt_id in specific_history, "History v2 should contain prompt ID" - + # Verify key fields match between legacy and v2 v2_data = specific_history[prompt_id] legacy_data = legacy_history[prompt_id] - + # Check that outputs and status match assert v2_data["outputs"] == legacy_data["outputs"], "Outputs should match" assert v2_data["status"] == legacy_data["status"], "Status should match" - + # Verify prompt is converted to dict format in v2 assert isinstance(v2_data["prompt"], dict), "Prompt should be a dictionary in v2" assert "prompt_id" in v2_data["prompt"], "Prompt dict should have prompt_id" @@ -761,18 +761,18 @@ class TestExecution: def test_ordered_history_prompt_field_filtering_unit(self): """Unit test for prompt field filtering logic in get_ordered_history.""" from execution import PromptQueue - + # Mock server class MockServer: def queue_updated(self): pass - + # Create queue and add mock history queue = PromptQueue(MockServer()) - + # Mock history entry with full prompt structure mock_prompt_tuple = ( 12345, # priority/timestamp - 'test-prompt-123', # prompt_id + 'test-prompt-123', # prompt_id {'nodes': {'1': {'class_type': 'SaveImage'}}}, # workflow (should be filtered) { # extra_data 'client_id': 'test-client', @@ -783,7 +783,7 @@ class TestExecution: }, ['1'] # execute_outputs (should be filtered) ) - + queue.history['test-prompt-123'] = { 'prompt': mock_prompt_tuple, 'outputs': {'1': {'images': []}}, @@ -797,53 +797,53 @@ class TestExecution: }, 'meta': {'1': {'node_id': '1'}} } - + # Test get_ordered_history with our filtering result = queue.get_ordered_history() - + # Verify structure assert "history" in result, "Result should have history key" assert len(result["history"]) == 1, "Should have one history item" - + history_item = result["history"][0] - + # Verify prompt_id field is added assert "prompt_id" in history_item, "History item should have prompt_id field" assert history_item["prompt_id"] == 'test-prompt-123', "prompt_id should match" - + # Verify prompt field is filtered filtered_prompt = history_item["prompt"] assert isinstance(filtered_prompt, dict), "Filtered prompt should be a dictionary" assert "priority" in filtered_prompt, "Filtered prompt should have priority" assert "prompt_id" in filtered_prompt, "Filtered prompt should have prompt_id" assert "extra_data" in filtered_prompt, "Filtered prompt should have extra_data" - + # Verify correct elements are preserved assert filtered_prompt["priority"] == 12345, "Priority should be preserved" assert filtered_prompt["prompt_id"] == 'test-prompt-123', "Prompt ID should be preserved" - + # Verify extra_data filtering extra_data = filtered_prompt["extra_data"] assert extra_data['client_id'] == 'test-client', "Client ID should be preserved" assert 'extra_pnginfo' in extra_data, "extra_pnginfo should be present" assert 'workflow' not in extra_data['extra_pnginfo'], "Workflow should be filtered out" assert extra_data['extra_pnginfo']['version'] == '1.0', "Other extra_pnginfo data should be preserved" - + # Verify other fields are unchanged assert history_item["outputs"] == {'1': {'images': []}}, "Outputs should be unchanged" assert history_item["meta"] == {'1': {'node_id': '1'}}, "Meta should be unchanged" - + # Verify status field filtering status = history_item["status"] assert "status_str" in status, "Status should have status_str" assert status["status_str"] == 'success', "Status string should be preserved" assert "completed" not in status, "Completed field should be filtered out" assert "messages" in status, "Status should have messages" - + # Verify message filtering messages = status["messages"] assert len(messages) == 2, "Should have 2 messages" - + # Check execution_cached message has nodes filtered out execution_cached_msg = messages[0] assert execution_cached_msg[0] == 'execution_cached', "First message should be execution_cached" @@ -851,7 +851,7 @@ class TestExecution: assert "nodes" not in cached_data, "Nodes field should be filtered from execution_cached messages" assert "timestamp" in cached_data, "Timestamp should be preserved in execution_cached messages" assert cached_data["timestamp"] == 1234567890, "Timestamp value should be correct" - + # Check execution_start message remains unchanged execution_start_msg = messages[1] assert execution_start_msg[0] == 'execution_start', "Second message should be execution_start"