Abineshsdata commited on
Commit
a1e9300
·
unverified ·
1 Parent(s): ffe399f

Implement tests for inference helper functions

Browse files

Add unit tests for inference helpers including action parsing and fallback policies.

Files changed (1) hide show
  1. tests/test_inference.py +181 -0
tests/test_inference.py ADDED
@@ -0,0 +1,181 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Test inference helpers.
3
+
4
+ These tests focus on the local Ollama/DeepSeek path where the model often
5
+ wraps JSON in reasoning text or code fences.
6
+ """
7
+
8
+ import sys
9
+ import os
10
+
11
+ sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
12
+
13
+ from inference import (
14
+ _is_ollama_backend,
15
+ _should_replace_with_fallback,
16
+ clamp_submission_score,
17
+ get_fallback_action,
18
+ parse_action,
19
+ )
20
+
21
+
22
+ class TestParseAction:
23
+ def test_parses_reasoning_plus_fenced_json(self):
24
+ response = """
25
+ <think>
26
+ I should bring up the hydro dam first.
27
+ </think>
28
+ ```json
29
+ {"action_type":"dispatch_generation","node_id":"NODE_01","mw":600}
30
+ ```
31
+ """
32
+ action = parse_action(response)
33
+ assert action == {
34
+ "action_type": "dispatch_generation",
35
+ "node_id": "NODE_01",
36
+ "mw": 600,
37
+ }
38
+
39
+ def test_parses_single_quoted_payload(self):
40
+ response = "{'action_type': 'toggle_circuit_breaker', 'edge_id': 'LINE_28', 'status': 'closed'}"
41
+ action = parse_action(response)
42
+ assert action == {
43
+ "action_type": "toggle_circuit_breaker",
44
+ "edge_id": "LINE_28",
45
+ "status": "CLOSED",
46
+ }
47
+
48
+ def test_parses_action_wrapped_in_list(self):
49
+ response = """
50
+ Here is the action:
51
+ [
52
+ {"action_type":"advance_tick"}
53
+ ]
54
+ """
55
+ action = parse_action(response)
56
+ assert action == {"action_type": "advance_tick"}
57
+
58
+
59
+ class TestClampSubmissionScore:
60
+ def test_clamp_keeps_scores_inside_open_interval(self):
61
+ assert 0.0 < clamp_submission_score(0.0) < 1.0
62
+ assert 0.0 < clamp_submission_score(1.0) < 1.0
63
+ assert clamp_submission_score(0.42) == 0.42
64
+
65
+
66
+ class TestOllamaDetection:
67
+ def test_local_ollama_url_detected_automatically(self):
68
+ assert _is_ollama_backend("http://localhost:11434/v1") is True
69
+
70
+ def test_remote_openai_compatible_url_requires_override(self):
71
+ assert _is_ollama_backend("https://example.com/v1") is False
72
+ assert _is_ollama_backend("https://example.com/v1", compat_override="1") is True
73
+
74
+
75
+ class TestFallbackPolicy:
76
+ def test_cascade_overload_replaces_non_isolation_action_while_line_is_overloaded(self):
77
+ obs = {
78
+ "topology_graph": {
79
+ "nodes": [
80
+ {"id": "NODE_01", "energized": True, "capacity_mw": 1500},
81
+ ],
82
+ "edges": [
83
+ {
84
+ "id": "LINE_29",
85
+ "status": "LIVE",
86
+ "current_load_mw": 950.0,
87
+ "capacity_mw": 1000.0,
88
+ },
89
+ ],
90
+ },
91
+ "telemetry_stream": [[
92
+ {"node_id": "NODE_01", "generation_mw": 600.0},
93
+ ]],
94
+ }
95
+
96
+ action = {"action_type": "dispatch_generation", "node_id": "NODE_01", "mw": 100}
97
+ assert _should_replace_with_fallback(task_id=2, tick=0, action_dict=action, obs_dict=obs)
98
+
99
+ def test_task3_replaces_early_dispatch_before_logs(self):
100
+ obs = {
101
+ "network_packet_logs": [],
102
+ "metadata": {"active_spoofs": ["NODE_14"]},
103
+ }
104
+ action = {"action_type": "dispatch_generation", "node_id": "NODE_09", "mw": 100}
105
+ assert _should_replace_with_fallback(task_id=3, tick=0, action_dict=action, obs_dict=obs)
106
+
107
+ def test_task3_replaces_early_quarantine_before_estimation(self):
108
+ obs = {
109
+ "network_packet_logs": [{"source_node": "NODE_14", "anomaly_flag": True}],
110
+ "metadata": {"active_spoofs": ["NODE_14"]},
111
+ }
112
+ action = {"action_type": "quarantine_scada_node", "node_id": "NODE_14"}
113
+ assert _should_replace_with_fallback(task_id=3, tick=1, action_dict=action, obs_dict=obs)
114
+
115
+ def test_task3_replaces_dispatch_before_quarantine(self):
116
+ obs = {
117
+ "network_packet_logs": [{"source_node": "NODE_14", "anomaly_flag": True}],
118
+ "last_state_estimation": {"consistent": False, "violation_node": "NODE_14"},
119
+ "metadata": {"active_spoofs": ["NODE_14"]},
120
+ }
121
+ action = {"action_type": "dispatch_generation", "node_id": "NODE_09", "mw": 100}
122
+ assert _should_replace_with_fallback(task_id=3, tick=2, action_dict=action, obs_dict=obs)
123
+
124
+ def test_task3_allows_reroute_after_quarantine(self):
125
+ obs = {
126
+ "topology_graph": {
127
+ "nodes": [
128
+ {"id": "NODE_09", "energized": True, "capacity_mw": 900},
129
+ ],
130
+ "edges": [],
131
+ },
132
+ "telemetry_stream": [[
133
+ {"node_id": "NODE_09", "generation_mw": 400.0},
134
+ ]],
135
+ "network_packet_logs": [{"source_node": "NODE_14", "anomaly_flag": True}],
136
+ "last_state_estimation": {"consistent": False, "violation_node": "NODE_14"},
137
+ "metadata": {"active_spoofs": []},
138
+ }
139
+ action = {"action_type": "dispatch_generation", "node_id": "NODE_09", "mw": 100}
140
+ assert _should_replace_with_fallback(task_id=3, tick=3, action_dict=action, obs_dict=obs) is False
141
+
142
+ def test_task4_replaces_non_injection_opening_move(self):
143
+ obs = {
144
+ "topology_graph": {"nodes": [], "edges": []},
145
+ "telemetry_stream": [],
146
+ }
147
+ action = {"action_type": "run_state_estimation", "subgraph": ["NODE_17", "NODE_20"]}
148
+ assert _should_replace_with_fallback(task_id=4, tick=0, action_dict=action, obs_dict=obs)
149
+
150
+ def test_task3_replaces_redundant_estimation_after_quarantine(self):
151
+ obs = {
152
+ "network_packet_logs": [{"source_node": "NODE_14", "anomaly_flag": True}],
153
+ "last_state_estimation": {"consistent": False, "violation_node": "NODE_14"},
154
+ "metadata": {"active_spoofs": []},
155
+ }
156
+ action = {"action_type": "run_state_estimation", "subgraph": ["NODE_14", "NODE_15"]}
157
+ assert _should_replace_with_fallback(task_id=3, tick=4, action_dict=action, obs_dict=obs)
158
+
159
+ def test_black_start_fallback_only_makes_minimal_progress(self):
160
+ obs = {
161
+ "topology_graph": {
162
+ "nodes": [
163
+ {"id": "NODE_01", "energized": True},
164
+ {"id": "NODE_03", "energized": False},
165
+ ],
166
+ "edges": [
167
+ {"id": "LINE_02", "status": "TRIPPED"},
168
+ ],
169
+ },
170
+ "telemetry_stream": [[
171
+ {"node_id": "NODE_01", "generation_mw": 600.0},
172
+ {"node_id": "NODE_03", "generation_mw": 0.0},
173
+ ]],
174
+ }
175
+
176
+ action = get_fallback_action(task_id=5, tick=3, obs_dict=obs)
177
+ assert action == {
178
+ "action_type": "toggle_circuit_breaker",
179
+ "edge_id": "LINE_02",
180
+ "status": "CLOSED",
181
+ }