Create infra_simulator.py

infra_simulator.py

@@ -1,34 +1,74 @@
-# infra_simulator.py
+"""
+Infrastructure simulator for data center topology with fault injection.
+Maintains per‑session state (to be stored in gr.State).
+"""
 import random
 from typing import Dict, List, Optional
 
 class InfraSimulator:
-    """Simulates a dynamic infrastructure. Maintains per‑session state."""
+    """Simulates a dynamic infrastructure with servers, switches, and services."""
+    
     def __init__(self):
         self.reset()
-
+    
     def reset(self):
+        """Reset to healthy state."""
         self.components = {
-            "switch-1": {"type": "switch", "status": "up", "connections": ["server-1", "server-2"]},
-            "server-1": {"type": "server", "status": "up", "connections": ["switch-1"], "services": ["db"]},
-            "server-2": {"type": "server", "status": "up", "connections": ["switch-1"], "services": ["web"]},
-            "service-db": {"type": "service", "status": "up", "runs_on": "server-1"},
-            "service-web": {"type": "service", "status": "up", "runs_on": "server-2"},
+            "switch-1": {
+                "type": "switch",
+                "status": "up",
+                "connections": ["server-1", "server-2"]
+            },
+            "server-1": {
+                "type": "server",
+                "status": "up",
+                "connections": ["switch-1"],
+                "services": ["db"]
+            },
+            "server-2": {
+                "type": "server",
+                "status": "up",
+                "connections": ["switch-1"],
+                "services": ["web"]
+            },
+            "service-db": {
+                "type": "service",
+                "status": "up",
+                "runs_on": "server-1"
+            },
+            "service-web": {
+                "type": "service",
+                "status": "up",
+                "runs_on": "server-2"
+            }
         }
         self.fault_mode: Optional[str] = None
-
+    
     def set_fault(self, fault: str):
-        self.reset()  # start fresh
+        """
+        Inject a fault.
+        Options: 'none', 'switch_down', 'server_overload', 'cascade'
+        """
+        self.reset()  # start from healthy
         self.fault_mode = fault
+        
         if fault == "switch_down":
             self.components["switch-1"]["status"] = "down"
+            # Propagate to connected servers (optional)
+            for conn in self.components["switch-1"]["connections"]:
+                if conn in self.components:
+                    self.components[conn]["status"] = "degraded"
         elif fault == "server_overload":
             self.components["server-1"]["status"] = "overloaded"
+            self.components["service-db"]["status"] = "degraded"
         elif fault == "cascade":
             self.components["switch-1"]["status"] = "down"
             self.components["server-1"]["status"] = "down"
+            self.components["server-2"]["status"] = "down"
             self.components["service-db"]["status"] = "down"
-        # etc.
-
+            self.components["service-web"]["status"] = "down"
+        # 'none' does nothing
+    
     def read_state(self) -> Dict:
+        """Return current component states."""
         return self.components