Skip to content

Commit

Permalink
make it a REST server with fastapi
Browse files Browse the repository at this point in the history
  • Loading branch information
YufengXin committed Nov 28, 2024
1 parent 9824e04 commit 2346b1f
Showing 2 changed files with 130 additions and 32 deletions.
156 changes: 126 additions & 30 deletions src/sdx_datamodel/connection_sm.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
import logging
from enum import Enum, auto

import matplotlib.pyplot as plt
import networkx as nx
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import FileResponse
from transitions import Machine
from transitions.extensions import GraphMachine
from enum import Enum, auto


class ConnectionStateMachine:
@@ -44,19 +46,71 @@ def __str__(self):
return self.name

transitions = [
{'trigger': str(Trigger.PROVISION), 'source': str(State.REQUESTED), 'dest': str(State.PROVISIONING)},
{'trigger': str(Trigger.REJECT), 'source': str(State.REQUESTED), 'dest': str(State.REJECTED)},
{'trigger': str(Trigger.PROVISION_SUCCESS), 'source': str(State.PROVISIONING), 'dest': str(State.PROVISIONED)},
{'trigger': str(Trigger.PROVISION_FAIL), 'source': str(State.PROVISIONING), 'dest': str(State.PROVISION_FAILED)},
{'trigger': str(Trigger.MODIFY), 'source': str(State.PROVISIONED), 'dest': str(State.MODIFYING)},
{'trigger': str(Trigger.MOD_SUCCESS), 'source': str(State.MODIFYING), 'dest': str(State.PROVISIONED)},
{'trigger': str(Trigger.MOD_FAIL), 'source': str(State.MODIFYING), 'dest': str(State.PROVISION_FAILED)},
{'trigger': str(Trigger.FAIL), 'source': str(State.PROVISIONED), 'dest': str(State.FAILED)},
{'trigger': str(Trigger.RECOVER), 'source': str(State.FAILED), 'dest': str(State.RECOVERING)},
{'trigger': str(Trigger.RECOVER_SUCCESS), 'source': str(State.RECOVERING), 'dest': str(State.PROVISIONED)},
{'trigger': str(Trigger.RECOVER_FAIL), 'source': str(State.RECOVERING), 'dest': str(State.FAILED)},
{'trigger': str(Trigger.DELETE), 'source': str(State.PROVISIONED), 'dest': str(State.DELETED)},
{'trigger': str(Trigger.DELETE), 'source': str(State.FAILED), 'dest': str(State.DELETED)}
{
"trigger": str(Trigger.PROVISION),
"source": str(State.REQUESTED),
"dest": str(State.PROVISIONING),
},
{
"trigger": str(Trigger.REJECT),
"source": str(State.REQUESTED),
"dest": str(State.REJECTED),
},
{
"trigger": str(Trigger.PROVISION_SUCCESS),
"source": str(State.PROVISIONING),
"dest": str(State.PROVISIONED),
},
{
"trigger": str(Trigger.PROVISION_FAIL),
"source": str(State.PROVISIONING),
"dest": str(State.PROVISION_FAILED),
},
{
"trigger": str(Trigger.MODIFY),
"source": str(State.PROVISIONED),
"dest": str(State.MODIFYING),
},
{
"trigger": str(Trigger.MOD_SUCCESS),
"source": str(State.MODIFYING),
"dest": str(State.PROVISIONED),
},
{
"trigger": str(Trigger.MOD_FAIL),
"source": str(State.MODIFYING),
"dest": str(State.PROVISION_FAILED),
},
{
"trigger": str(Trigger.FAIL),
"source": str(State.PROVISIONED),
"dest": str(State.FAILED),
},
{
"trigger": str(Trigger.RECOVER),
"source": str(State.FAILED),
"dest": str(State.RECOVERING),
},
{
"trigger": str(Trigger.RECOVER_SUCCESS),
"source": str(State.RECOVERING),
"dest": str(State.PROVISIONED),
},
{
"trigger": str(Trigger.RECOVER_FAIL),
"source": str(State.RECOVERING),
"dest": str(State.FAILED),
},
{
"trigger": str(Trigger.DELETE),
"source": str(State.PROVISIONED),
"dest": str(State.DELETED),
},
{
"trigger": str(Trigger.DELETE),
"source": str(State.FAILED),
"dest": str(State.DELETED),
},
]

def __init__(self):
@@ -65,28 +119,34 @@ def __init__(self):

def transition(self, new_state):
valid_transitions = {

self.State.REQUESTED: [self.State.PROVISIONING],

self.State.PROVISIONING: [self.State.PROVISIONED, self.State.PROVISION_FAILED],

self.State.PROVISIONED: [self.State.MODIFYING, self.State.FAILED, self.State.DELETED],

self.State.PROVISIONING: [
self.State.PROVISIONED,
self.State.PROVISION_FAILED,
],
self.State.PROVISIONED: [
self.State.MODIFYING,
self.State.FAILED,
self.State.DELETED,
],
self.State.PROVISION_FAILED: [self.State.RECOVERING],

self.State.MODIFYING: [self.State.PROVISIONED, self.State.PROVISION_FAILED],

self.State.MODIFYING: [
self.State.PROVISIONED,
self.State.PROVISION_FAILED,
],
self.State.FAILED: [self.State.RECOVERING],

self.State.RECOVERING: [self.State.PROVISIONED, self.State.PROVISION_FAILED],

self.State.DELETED: []

self.State.RECOVERING: [
self.State.PROVISIONED,
self.State.PROVISION_FAILED,
],
self.State.DELETED: [],
}

if new_state not in valid_transitions[self.state]:

raise ValueError(f"Invalid transition from {self.state} to {new_state}")
raise ValueError(
f"Invalid transition from {self.state} to {new_state}"
)

self.state = new_state

@@ -102,7 +162,8 @@ def __str__(self):
def set_state(self, state):
self.state = state

def draw_transition(model,output):

def draw_transition(model, output):
machine = GraphMachine(
model=model,
graph_engine="graphviz",
@@ -118,6 +179,41 @@ def draw_transition(model,output):
# you need to pass the format when you pass objects instead of filenames.
machine.get_graph().draw(f, format="png", prog="dot")


app = FastAPI()
connections = {}


@app.get("/")
async def root():
return {"message": "Connection State Machine Server"}


@app.post("/connection/")
def create_connection():
connection_id = len(connections) + 1
connection = ConnectionStateMachine()
connections[connection_id] = connection
return {"connection_id": connection_id, "state": connection.state}


@app.put("/connections/{connection_id}/provision/")
def process_connection(connection_id: int):
connection = connections.get(connection_id)
if connection:
connection.transition(connection.State.PROVISIONING)
return {"connection_id": connection_id, "state": connection.state}
return {"error": "connection not found"}


@app.get("/connection/sm/")
def get_connection_state_machine():
connection = ConnectionStateMachine()
draw_transition(connection, "connection_transition.png")
# return {"message": "Connection State Machine"}
return FileResponse("./connection_transition.png")


if __name__ == "__main__":
sm = ConnectionStateMachine()
draw_transition(sm, "connection_transition.png")
6 changes: 4 additions & 2 deletions tests/test_connection_sm.py
Original file line number Diff line number Diff line change
@@ -2,6 +2,7 @@

from src.sdx_datamodel.connection_sm import ConnectionStateMachine


class TestConnectionStateMachine(unittest.TestCase):

def setUp(self):
@@ -61,5 +62,6 @@ def test_delete(self):
self.sm.transition(self.State.DELETED)
self.assertEqual(self.sm.get_state(), self.State.DELETED)

if __name__ == '__main__':
unittest.main()

if __name__ == "__main__":
unittest.main()

0 comments on commit 2346b1f

Please sign in to comment.