Support arg_to_param_map (dict instead of function)

2 files changed, 29 insertions, 4 deletions

diff --git a/bin/test_automata.py b/bin/test_automata.py
index 5a62931..c21a5a1 100755
--- a/bin/test_automata.py
+++ b/bin/test_automata.py
@@ -152,7 +152,7 @@ class TestPTA(unittest.TestCase):
             'length' : None
         })
 
-    def test_simulation_param_set(self):
+    def test_simulation_param_update_function(self):
         pta = PTA(parameters = ['txpower', 'length'])
         pta.add_state('IDLE', power = 5)
         pta.add_state('TX', power = 100)
@@ -176,6 +176,30 @@ class TestPTA(unittest.TestCase):
             'length' : None
         })
 
+    def test_simulation_arg_to_param_map(self):
+        pta = PTA(parameters = ['txpower', 'length'])
+        pta.add_state('IDLE', power = 5)
+        pta.add_state('TX', power = 100)
+        pta.add_transition('UNINITIALIZED', 'IDLE', 'init', energy = 500000, duration = 50000)
+        pta.add_transition('IDLE', 'IDLE', 'setTxPower', energy = 10000, duration = 120,
+            arg_to_param_map = {'txpower' : 0})
+        pta.add_transition('IDLE', 'TX', 'send', energy = 3, duration = 10)
+        pta.add_transition('TX', 'IDLE', 'txComplete', timeout = 2000, is_interrupt = True)
+        trace = [
+            ['init'],
+            ['setTxPower', 10]
+        ]
+        expected_energy = 510000
+        expected_duration = 50120
+        power, duration, state, parameters = pta.simulate(trace)
+        self.assertEqual(power, expected_energy)
+        self.assertEqual(duration, expected_duration)
+        self.assertEqual(state.name, 'IDLE')
+        self.assertEqual(parameters, {
+            'txpower' : 10,
+            'length' : None
+        })
+
     def test_simulation_arg_function(self):
         pta = PTA(parameters = ['txpower', 'length'])
         pta.add_state('IDLE', power = 5)
diff --git a/lib/automata.py b/lib/automata.py
index 0948cee..ff32bf7 100755
--- a/lib/automata.py
+++ b/lib/automata.py
@@ -51,9 +51,10 @@ class Transition:
         if self.param_update_function:
             return self.param_update_function(param_dict, args)
         if self.arg_to_param_map:
-            ret = {}
-            #for arg_index in range(self.arg_to_param_map):
-            #    if self.arg_to_param_map[arg_index]
+            ret = param_dict.copy()
+            for k, v in self.arg_to_param_map.items():
+                ret[k] = args[v]
+            return ret
         return param_dict
 
 class State: