updated test file

src/environment.py

@@ -96,9 +96,9 @@ class AdPolicyEnvironment(Environment):
                 return self._get_obs("Must call query_regulations first.", -0.2, False)
 
         self.step_count += 1
-        self.actions_taken.add(act_type)
+        
 
-        # 3. Execute action
+       # 3. Execute action
         response = self._execute_action(act_type)
 
         # 4. Update state
@@ -106,7 +106,9 @@ class AdPolicyEnvironment(Environment):
             self.api_failed = True
             self.last_failed_action = act_type
             self.last_error = response["error"]
+            # Notice we DO NOT add to self.actions_taken here
         else:
+            self.actions_taken.add(act_type)  # <-- ADDED HERE: Only register successful actions
             if act_type == self.last_failed_action:
                 self.api_recovered = True
             self.last_error = None

test_env.py

@@ -1,57 +1,65 @@
-import requests
-import json
-
-BASE_URL = "http://localhost:8000"
-
-def safe_post(endpoint, data=None):
-    """Helper to catch errors before they crash the script."""
-    try:
-        url = f"{BASE_URL}/{endpoint}"
-        response = requests.post(url, json=data)
-        
-        # If the server sent an error code (4xx or 5xx), print the text
-        if response.status_code != 200:
-            print(f"❌ Server Error {response.status_code}: {response.text}")
-            return None
-            
-        return response.json()
-    except Exception as e:
-        print(f"⚠️ Request Failed: {e}")
-        return None
-
-def run_test():
-    print("--- 🔄 Testing /reset ---")
-    reset_data = safe_post("reset")
-    if not reset_data: return
-    
-    obs = reset_data.get('observation', reset_data)
-    print(f"Ad Loaded: {obs.get('headline', 'N/A')}\n")
-
-    print("--- 🔍 Testing 'analyze_image' Tool ---")
-    # Payload must be wrapped in 'action' for OpenEnv 2026
-    step1_payload = {
-        "action": {
-            "action_type": "analyze_image",
-            "reasoning": "Standard adversarial check."
-        }
-    }
-    s1_data = safe_post("step", step1_payload)
-    if s1_data:
-        s1_obs = s1_data.get('observation', s1_data)
-        print(f" {s1_obs.get('status_message', 'N/A')}\n")
-
-    print("--- ✅ Testing Final Decision ---")
-    step2_payload = {
-        "action": {
-            "action_type": "reject",
-            "reasoning": "Detected policy violation."
-        }
-    }
-    s2_data = safe_post("step", step2_payload)
-    if s2_data:
-        reward = s2_data.get('reward', 0.0)
-        done = s2_data.get('done', s2_data.get('terminal', False))
-        print(f"Final Reward: {reward} | Done: {done}")
+from src.environment import AdPolicyEnvironment
+from src.models import AdAction
+
+
+def run_episode(task_id, actions):
+    env = AdPolicyEnvironment()
+    obs = env.reset(task_id=task_id)
+
+    for act in actions:
+        obs = env.step(
+            AdAction(
+                action_type=act,
+                reasoning="smoke test",
+                violation_category="NONE",
+            )
+        )
+        if obs.done:
+            break
+
+    return env, obs
+
 
 if __name__ == "__main__":
-    run_test()
+    env1, obs1 = run_episode(
+        "task_1_healthcare",
+        [
+            "query_regulations",
+            "analyze_image",
+            "check_advertiser_history",
+            "submit_audit",
+            "reject",
+        ],
+    )
+
+    assert len(env1.trace) >= 4, f"Trace too short: {len(env1.trace)}"
+    assert isinstance(env1.total_reward, float), "Reward is not numeric"
+    assert all("summary" in t["result"] for t in env1.trace), "Bad trace format"
+
+    env2, obs2 = run_episode(
+        "task_10_failure",
+        [
+            "query_regulations",
+            "query_regulations",
+            "check_advertiser_history",
+            "submit_audit",
+            "reject",
+        ],
+    )
+
+    assert len(env2.trace) >= 2, f"Failure trace too short: {len(env2.trace)}"
+    assert any("API failure" in t["result"]["summary"] for t in env2.trace), (
+        "Failure case did not trigger"
+    )
+
+    print("STEP 7 SMOKE TEST PASSED")
+    print("\nTRACE 1:")
+    for row in env1.trace:
+        print(row)
+
+    print("\nTRACE 2:")
+    for row in env2.trace:
+        print(row)
+
+    print("\nTOTAL REWARD 1:", env1.total_reward)
+    print("TOTAL REWARD 2:", env2.total_reward)