autopep8

1 files changed, 240 insertions, 243 deletions

diff --git a/test/test_pta.py b/test/test_pta.py
index ff67fb5..5158005 100755
--- a/test/test_pta.py
+++ b/test/test_pta.py
@@ -5,81 +5,81 @@ import unittest
 import yaml
 
 example_json_1 = {
-    'parameters' : ['datarate', 'txbytes', 'txpower'],
-    'initial_param_values' : [None, None, None],
-    'state' : {
-        'IDLE' : {
-            'power' : {
-                'static' : 5,
+    'parameters': ['datarate', 'txbytes', 'txpower'],
+    'initial_param_values': [None, None, None],
+    'state': {
+        'IDLE': {
+            'power': {
+                'static': 5,
             }
         },
-        'TX' : {
-            'power' : {
-                'static' : 10000,
-                'function' : {
-                    'raw' : 'regression_arg(0) + regression_arg(1)'
-                        ' * parameter(txpower)',
-                    'regression_args' : [ 10000, 2 ]
+        'TX': {
+            'power': {
+                'static': 10000,
+                'function': {
+                    'raw': 'regression_arg(0) + regression_arg(1)'
+                    ' * parameter(txpower)',
+                    'regression_args': [10000, 2]
                 },
             }
         },
     },
-    'transitions' : [
+    'transitions': [
         {
-            'name' : 'init',
-            'origin' : ['UNINITIALIZED', 'IDLE'],
-            'destination' : 'IDLE',
-            'duration' : {
-                'static' : 50000,
+            'name': 'init',
+            'origin': ['UNINITIALIZED', 'IDLE'],
+            'destination': 'IDLE',
+            'duration': {
+                'static': 50000,
             },
-            'set_param' : {
-                'txpower' : 10
+            'set_param': {
+                'txpower': 10
             },
         },
         {
-            'name' : 'setTxPower',
-            'origin' : 'IDLE',
-            'destination' : 'IDLE',
-            'duration' : { 'static' : 120 },
-            'energy ' : { 'static' : 10000 },
-            'arg_to_param_map' : { 0: 'txpower'},
-            'argument_values' : [ [10, 20, 30] ],
+            'name': 'setTxPower',
+            'origin': 'IDLE',
+            'destination': 'IDLE',
+            'duration': {'static': 120},
+            'energy ': {'static': 10000},
+            'arg_to_param_map': {0: 'txpower'},
+            'argument_values': [[10, 20, 30]],
         },
         {
-            'name' : 'send',
-            'origin' : 'IDLE',
-            'destination' : 'TX',
-            'duration' : {
-                'static' : 10,
-                'function' : {
-                    'raw' : 'regression_arg(0) + regression_arg(1)'
-                        ' * function_arg(1)',
-                    'regression_args' : [48, 8],
+            'name': 'send',
+            'origin': 'IDLE',
+            'destination': 'TX',
+            'duration': {
+                'static': 10,
+                'function': {
+                    'raw': 'regression_arg(0) + regression_arg(1)'
+                    ' * function_arg(1)',
+                    'regression_args': [48, 8],
                 },
             },
-            'energy' : {
-                'static' : 3,
-                'function' : {
-                    'raw' : 'regression_arg(0) + regression_arg(1)'
-                        ' * function_arg(1)',
-                    'regression_args' : [3, 5],
+            'energy': {
+                'static': 3,
+                'function': {
+                    'raw': 'regression_arg(0) + regression_arg(1)'
+                    ' * function_arg(1)',
+                    'regression_args': [3, 5],
                 },
             },
-            'arg_to_param_map' : { 1: 'txbytes'},
-            'argument_values' : [ ['"foo"', '"hodor"'], [3, 5] ],
-            'argument_combination' : 'zip',
+            'arg_to_param_map': {1: 'txbytes'},
+            'argument_values': [['"foo"', '"hodor"'], [3, 5]],
+            'argument_combination': 'zip',
         },
         {
-            'name' : 'txComplete',
-            'origin' : 'TX',
-            'destination' : 'IDLE',
-            'is_interrupt' : 1,
-            'timeout' : {
-                'static' : 2000,
-                'function' : {
-                    'raw' : 'regression_arg(0) + regression_arg(1)'
-                        ' * parameter(txbytes)',
-                    'regression_args' : [ 500, 16 ],
+            'name': 'txComplete',
+            'origin': 'TX',
+            'destination': 'IDLE',
+            'is_interrupt': 1,
+            'timeout': {
+                'static': 2000,
+                'function': {
+                    'raw': 'regression_arg(0) + regression_arg(1)'
+                    ' * parameter(txbytes)',
+                    'regression_args': [500, 16],
                 },
             },
         }
@@ -262,6 +262,7 @@ transition:
     argument_combination: zip
 """)
 
+
 def dfs_tran_to_name(runs: list, with_args: bool = False, with_param: bool = False) -> list:
     if with_param:
         return list(map(lambda run: list(map(lambda x: (x[0].name, x[1], x[2]), run)), runs))
@@ -269,6 +270,7 @@ def dfs_tran_to_name(runs: list, with_args: bool = False, with_param: bool = Fal
         return list(map(lambda run: list(map(lambda x: (x[0].name, x[1]), run)), runs))
     return list(map(lambda run: list(map(lambda x: (x[0].name), run)), runs))
 
+
 class TestPTA(unittest.TestCase):
     def test_dfs(self):
         pta = PTA(['IDLE', 'TX'])
@@ -287,9 +289,9 @@ class TestPTA(unittest.TestCase):
         self.assertEqual(dfs_tran_to_name(pta.dfs(0), False), [['init']])
         self.assertEqual(sorted(dfs_tran_to_name(pta.dfs(1), False)), [['init', 'set1'], ['init', 'set2']])
         self.assertEqual(sorted(dfs_tran_to_name(pta.dfs(2), False)), [['init', 'set1', 'set1'],
-            ['init', 'set1', 'set2'],
-            ['init', 'set2', 'set1'],
-            ['init', 'set2', 'set2']])
+                                                                       ['init', 'set1', 'set2'],
+                                                                       ['init', 'set2', 'set1'],
+                                                                       ['init', 'set2', 'set2']])
 
     def test_dfs_trace_filter(self):
         pta = PTA(['IDLE'])
@@ -297,12 +299,12 @@ class TestPTA(unittest.TestCase):
         pta.add_transition('IDLE', 'IDLE', 'set1')
         pta.add_transition('IDLE', 'IDLE', 'set2')
         self.assertEqual(sorted(dfs_tran_to_name(pta.dfs(2, trace_filter=[['init', 'set1', 'set2'], ['init', 'set2', 'set1']]), False)),
-            [['init', 'set1', 'set2'], ['init', 'set2', 'set1']])
+                         [['init', 'set1', 'set2'], ['init', 'set2', 'set1']])
         self.assertEqual(sorted(dfs_tran_to_name(pta.dfs(2, trace_filter=[['init', 'set1', '$'], ['init', 'set2', '$']]), False)),
-            [['init', 'set1'], ['init', 'set2']])
+                         [['init', 'set1'], ['init', 'set2']])
 
     def test_dfs_accepting(self):
-        pta = PTA(['IDLE', 'TX'], accepting_states = ['IDLE'])
+        pta = PTA(['IDLE', 'TX'], accepting_states=['IDLE'])
         pta.add_transition('UNINITIALIZED', 'IDLE', 'init')
         pta.add_transition('IDLE', 'TX', 'send')
         pta.add_transition('TX', 'IDLE', 'txComplete')
@@ -332,7 +334,7 @@ class TestPTA(unittest.TestCase):
         pta.add_transition('UNINITIALIZED', 'IDLE', 'init')
         pta.add_transition('IDLE', 'TX', 'send')
         pta.add_transition('TX', 'IDLE', 'txComplete')
-        traces = list(pta.dfs(2, sleep = 10))
+        traces = list(pta.dfs(2, sleep=10))
         self.assertEqual(len(traces), 1)
         trace = traces[0]
         self.assertEqual(len(trace), 6)
@@ -359,7 +361,7 @@ class TestPTA(unittest.TestCase):
         self.assertEqual(pta.transitions[3].name, 'send')
         self.assertEqual(pta.transitions[4].name, 'txComplete')
 
-    #def test_to_json(self):
+    # def test_to_json(self):
     #    pta = PTA.from_json(example_json_1)
     #    json = pta.to_json()
     #    json['state'].pop('UNINITIALIZED')
@@ -369,152 +371,149 @@ class TestPTA(unittest.TestCase):
     def test_from_json_dfs_arg(self):
         pta = PTA.from_json(example_json_1)
         self.assertEqual(sorted(dfs_tran_to_name(pta.dfs(1), False)), [['init', 'init'], ['init', 'send'], ['init', 'setTxPower']])
-        self.assertEqual(sorted(dfs_tran_to_name(pta.dfs(1, with_arguments = True), True)),
-            [
-                [('init', ()), ('init', ())],
-                [('init', ()), ('send', ('"foo"', 3))],
-                [('init', ()), ('send', ('"hodor"', 5))],
-                [('init', ()), ('setTxPower', (10,))],
-                [('init', ()), ('setTxPower', (20,))],
-                [('init', ()), ('setTxPower', (30,))],
-            ]
+        self.assertEqual(sorted(dfs_tran_to_name(pta.dfs(1, with_arguments=True), True)),
+                         [
+            [('init', ()), ('init', ())],
+            [('init', ()), ('send', ('"foo"', 3))],
+            [('init', ()), ('send', ('"hodor"', 5))],
+            [('init', ()), ('setTxPower', (10,))],
+            [('init', ()), ('setTxPower', (20,))],
+            [('init', ()), ('setTxPower', (30,))],
+        ]
         )
 
     def test_from_json_dfs_param(self):
         pta = PTA.from_json(example_json_1)
         no_param = {
-            'datarate' : None,
-            'txbytes' : None,
-            'txpower' : None,
+            'datarate': None,
+            'txbytes': None,
+            'txpower': None,
         }
         param_tx3 = {
-            'datarate' : None,
-            'txbytes' : 3,
-            'txpower' : None,
+            'datarate': None,
+            'txbytes': 3,
+            'txpower': None,
         }
         param_tx5 = {
-            'datarate' : None,
-            'txbytes' : 5,
-            'txpower' : None,
+            'datarate': None,
+            'txbytes': 5,
+            'txpower': None,
         }
         param_txp10 = {
-            'datarate' : None,
-            'txbytes' : None,
-            'txpower' : 10,
+            'datarate': None,
+            'txbytes': None,
+            'txpower': 10,
         }
         param_txp20 = {
-            'datarate' : None,
-            'txbytes' : None,
-            'txpower' : 20,
+            'datarate': None,
+            'txbytes': None,
+            'txpower': 20,
         }
         param_txp30 = {
-            'datarate' : None,
-            'txbytes' : None,
-            'txpower' : 30,
+            'datarate': None,
+            'txbytes': None,
+            'txpower': 30,
         }
-        self.assertEqual(sorted(dfs_tran_to_name(pta.dfs(1, with_arguments = True, with_parameters = True), True, True)),
-            [
-                [('init', (), no_param), ('init', (), no_param)],
-                [('init', (), no_param), ('send', ('"foo"', 3), param_tx3)],
-                [('init', (), no_param), ('send', ('"hodor"', 5), param_tx5)],
-                [('init', (), no_param), ('setTxPower', (10,), param_txp10)],
-                [('init', (), no_param), ('setTxPower', (20,), param_txp20)],
-                [('init', (), no_param), ('setTxPower', (30,), param_txp30)],
-            ]
+        self.assertEqual(sorted(dfs_tran_to_name(pta.dfs(1, with_arguments=True, with_parameters=True), True, True)),
+                         [
+            [('init', (), no_param), ('init', (), no_param)],
+            [('init', (), no_param), ('send', ('"foo"', 3), param_tx3)],
+            [('init', (), no_param), ('send', ('"hodor"', 5), param_tx5)],
+            [('init', (), no_param), ('setTxPower', (10,), param_txp10)],
+            [('init', (), no_param), ('setTxPower', (20,), param_txp20)],
+            [('init', (), no_param), ('setTxPower', (30,), param_txp30)],
+        ]
         )
 
     def test_from_json_function(self):
         pta = PTA.from_json(example_json_1)
-        self.assertEqual(pta.state['TX'].get_energy(1000, {'datarate' : 10, 'txbytes' : 6, 'txpower' : 10 }), 1000 * (10000 + 2 * 10))
-        self.assertEqual(pta.transitions[4].get_timeout({'datarate' : 10, 'txbytes' : 6, 'txpower' : 10 }), 500 + 16 * 6)
-
-    def test_from_yaml(self):
-        pta = PTA.from_yaml(example_yaml_1)
+        self.assertEqual(pta.state['TX'].get_energy(1000, {'datarate': 10, 'txbytes': 6, 'txpower': 10}), 1000 * (10000 + 2 * 10))
+        self.assertEqual(pta.transitions[4].get_timeout({'datarate': 10, 'txbytes': 6, 'txpower': 10}), 500 + 16 * 6)
 
     def test_from_yaml_dfs_param(self):
         pta = PTA.from_yaml(example_yaml_1)
         no_param = {
-            'datarate' : None,
-            'txbytes' : None,
-            'txpower' : None,
+            'datarate': None,
+            'txbytes': None,
+            'txpower': None,
         }
         param_tx3 = {
-            'datarate' : None,
-            'txbytes' : 3,
-            'txpower' : None,
+            'datarate': None,
+            'txbytes': 3,
+            'txpower': None,
         }
         param_tx5 = {
-            'datarate' : None,
-            'txbytes' : 5,
-            'txpower' : None,
+            'datarate': None,
+            'txbytes': 5,
+            'txpower': None,
         }
         param_txp10 = {
-            'datarate' : None,
-            'txbytes' : None,
-            'txpower' : 10,
+            'datarate': None,
+            'txbytes': None,
+            'txpower': 10,
         }
         param_txp20 = {
-            'datarate' : None,
-            'txbytes' : None,
-            'txpower' : 20,
+            'datarate': None,
+            'txbytes': None,
+            'txpower': 20,
         }
         param_txp30 = {
-            'datarate' : None,
-            'txbytes' : None,
-            'txpower' : 30,
+            'datarate': None,
+            'txbytes': None,
+            'txpower': 30,
         }
-        self.assertEqual(sorted(dfs_tran_to_name(pta.dfs(1, with_arguments = True, with_parameters = True), True, True)),
-            [
-                [('init', (), no_param), ('init', (), no_param)],
-                [('init', (), no_param), ('send', ('"foo"', 3), param_tx3)],
-                [('init', (), no_param), ('send', ('"hodor"', 5), param_tx5)],
-                [('init', (), no_param), ('setTxPower', (10,), param_txp10)],
-                [('init', (), no_param), ('setTxPower', (20,), param_txp20)],
-                [('init', (), no_param), ('setTxPower', (30,), param_txp30)],
-            ]
+        self.assertEqual(sorted(dfs_tran_to_name(pta.dfs(1, with_arguments=True, with_parameters=True), True, True)),
+                         [
+            [('init', (), no_param), ('init', (), no_param)],
+            [('init', (), no_param), ('send', ('"foo"', 3), param_tx3)],
+            [('init', (), no_param), ('send', ('"hodor"', 5), param_tx5)],
+            [('init', (), no_param), ('setTxPower', (10,), param_txp10)],
+            [('init', (), no_param), ('setTxPower', (20,), param_txp20)],
+            [('init', (), no_param), ('setTxPower', (30,), param_txp30)],
+        ]
         )
 
     def test_normalization(self):
         pta = PTA.from_yaml(example_yaml_2)
         no_param = {
-            'datarate' : None,
-            'txbytes' : None,
-            'txpower' : None,
+            'datarate': None,
+            'txbytes': None,
+            'txpower': None,
         }
         param_tx3 = {
-            'datarate' : None,
-            'txbytes' : 3,
-            'txpower' : None,
+            'datarate': None,
+            'txbytes': 3,
+            'txpower': None,
         }
         param_tx6 = {
-            'datarate' : None,
-            'txbytes' : 6,
-            'txpower' : None,
+            'datarate': None,
+            'txbytes': 6,
+            'txpower': None,
         }
         param_txp10 = {
-            'datarate' : None,
-            'txbytes' : None,
-            'txpower' : -6,
+            'datarate': None,
+            'txbytes': None,
+            'txpower': -6,
         }
         param_txp20 = {
-            'datarate' : None,
-            'txbytes' : None,
-            'txpower' : 4,
+            'datarate': None,
+            'txbytes': None,
+            'txpower': 4,
         }
         param_txp30 = {
-            'datarate' : None,
-            'txbytes' : None,
-            'txpower' : 14,
+            'datarate': None,
+            'txbytes': None,
+            'txpower': 14,
         }
-        self.assertEqual(sorted(dfs_tran_to_name(pta.dfs(1, with_arguments = True, with_parameters = True), True, True)),
-            [
-                [('init', (), no_param), ('init', (), no_param)],
-                [('init', (), no_param), ('send', ('FOO',), param_tx3)],
-                [('init', (), no_param), ('send', ('LONGER',), param_tx6)],
-                [('init', (), no_param), ('setTxPower', (10,), param_txp10)],
-                [('init', (), no_param), ('setTxPower', (20,), param_txp20)],
-                [('init', (), no_param), ('setTxPower', (30,), param_txp30)],
-            ]
+        self.assertEqual(sorted(dfs_tran_to_name(pta.dfs(1, with_arguments=True, with_parameters=True), True, True)),
+                         [
+            [('init', (), no_param), ('init', (), no_param)],
+            [('init', (), no_param), ('send', ('FOO',), param_tx3)],
+            [('init', (), no_param), ('send', ('LONGER',), param_tx6)],
+            [('init', (), no_param), ('setTxPower', (10,), param_txp10)],
+            [('init', (), no_param), ('setTxPower', (20,), param_txp20)],
+            [('init', (), no_param), ('setTxPower', (30,), param_txp30)],
+        ]
         )
 
     def test_shrink(self):
@@ -530,15 +529,15 @@ class TestPTA(unittest.TestCase):
         self.assertEqual(pta.transitions[1].argument_values, [['Nrf24l01::RF24_PA_MIN', 'Nrf24l01::RF24_PA_MAX']])
         self.assertEqual(pta.transitions[2].argument_values, [[0, 15], [0, 15]])
         self.assertEqual(pta.transitions[5].argument_values, [['"foo"', '"foo"', '"foofoofoo"', '"foofoofoo"', '"123456789012345678901234567890"',
-        '"123456789012345678901234567890"'], [3, 3, 9, 9, 30, 30], [0, 1, 0, 1, 0, 1],  [1, 1, 1, 1, 1, 1]])
+                                                               '"123456789012345678901234567890"'], [3, 3, 9, 9, 30, 30], [0, 1, 0, 1, 0, 1], [1, 1, 1, 1, 1, 1]])
 
     def test_simulation(self):
         pta = PTA()
-        pta.add_state('IDLE', power = 5)
-        pta.add_state('TX', power = 100)
-        pta.add_transition('UNINITIALIZED', 'IDLE', 'init', duration = 50000)
-        pta.add_transition('IDLE', 'TX', 'send', energy = 3, duration = 10)
-        pta.add_transition('TX', 'IDLE', 'txComplete', timeout = 2000, is_interrupt = True)
+        pta.add_state('IDLE', power=5)
+        pta.add_state('TX', power=100)
+        pta.add_transition('UNINITIALIZED', 'IDLE', 'init', duration=50000)
+        pta.add_transition('IDLE', 'TX', 'send', energy=3, duration=10)
+        pta.add_transition('TX', 'IDLE', 'txComplete', timeout=2000, is_interrupt=True)
         trace = [
             ['init'],
             [None, 10000000],
@@ -555,12 +554,12 @@ class TestPTA(unittest.TestCase):
         self.assertEqual(parameters, {})
 
     def test_simulation_param_none(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', 'TX', 'send', energy = 3, duration = 10)
-        pta.add_transition('TX', 'IDLE', 'txComplete', timeout = 2000, is_interrupt = True)
+        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', 'TX', 'send', energy=3, duration=10)
+        pta.add_transition('TX', 'IDLE', 'txComplete', timeout=2000, is_interrupt=True)
         trace = [
             ['init'],
         ]
@@ -571,19 +570,19 @@ class TestPTA(unittest.TestCase):
         self.assertEqual(duration, expected_duration)
         self.assertEqual(state.name, 'IDLE')
         self.assertEqual(parameters, {
-            'txpower' : None,
-            'length' : None
+            'txpower': None,
+            'length': None
         })
 
     def test_simulation_param_update_function(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,
-            param_update_function = lambda param, arg: {**param, 'txpower' : arg[0]})
-        pta.add_transition('IDLE', 'TX', 'send', energy = 3, duration = 10)
-        pta.add_transition('TX', 'IDLE', 'txComplete', timeout = 2000, is_interrupt = True)
+        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,
+                           param_update_function=lambda param, arg: {**param, 'txpower': arg[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]
@@ -595,19 +594,19 @@ class TestPTA(unittest.TestCase):
         self.assertEqual(duration, expected_duration)
         self.assertEqual(state.name, 'IDLE')
         self.assertEqual(parameters, {
-            'txpower' : 10,
-            'length' : None
+            'txpower': 10,
+            '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 = {0: 'txpower'})
-        pta.add_transition('IDLE', 'TX', 'send', energy = 3, duration = 10)
-        pta.add_transition('TX', 'IDLE', 'txComplete', timeout = 2000, is_interrupt = True)
+        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={0: 'txpower'})
+        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]
@@ -619,15 +618,15 @@ class TestPTA(unittest.TestCase):
         self.assertEqual(duration, expected_duration)
         self.assertEqual(state.name, 'IDLE')
         self.assertEqual(parameters, {
-            'txpower' : 10,
-            'length' : None
+            'txpower': 10,
+            'length': None
         })
 
     def test_simulation_set_param(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, set_param = {'txpower' : 10})
+        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, set_param={'txpower': 10})
         trace = [
             ['init'],
         ]
@@ -638,21 +637,21 @@ class TestPTA(unittest.TestCase):
         self.assertEqual(duration, expected_duration)
         self.assertEqual(state.name, 'IDLE')
         self.assertEqual(parameters, {
-            'txpower' : 10,
-            'length' : None
+            'txpower': 10,
+            'length': None
         })
 
     def test_simulation_arg_function(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,
-            param_update_function = lambda param, arg: {**param, 'txpower' : arg[0]})
-        pta.add_transition('IDLE', 'TX', 'send', energy = 3, duration = 10,
-            energy_function = lambda param, arg: 3 + 5 * arg[1],
-            duration_function = lambda param, arg: 48 + 8 * arg[1])
-        pta.add_transition('TX', 'IDLE', 'txComplete', timeout = 2000, is_interrupt = True)
+        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,
+                           param_update_function=lambda param, arg: {**param, 'txpower': arg[0]})
+        pta.add_transition('IDLE', 'TX', 'send', energy=3, duration=10,
+                           energy_function=lambda param, arg: 3 + 5 * arg[1],
+                           duration_function=lambda param, arg: 48 + 8 * arg[1])
+        pta.add_transition('TX', 'IDLE', 'txComplete', timeout=2000, is_interrupt=True)
         trace = [
             ['init'],
             ['setTxPower', 10],
@@ -665,20 +664,20 @@ class TestPTA(unittest.TestCase):
         self.assertEqual(duration, expected_duration)
         self.assertEqual(state.name, 'IDLE')
         self.assertEqual(parameters, {
-            'txpower' : 10,
-            'length' : None
+            'txpower': 10,
+            'length': None
         })
 
-        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,
-            param_update_function = lambda param, arg: {**param, 'txpower' : arg[0]})
-        pta.add_transition('IDLE', 'TX', 'send', energy = 3, duration = 10,
-            energy_function = lambda param, arg: 3 + 5 * arg[1],
-            duration_function = lambda param, arg: 48 + 8 * arg[1])
-        pta.add_transition('TX', 'IDLE', 'txComplete', timeout = 2000, is_interrupt = True)
+        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,
+                           param_update_function=lambda param, arg: {**param, 'txpower': arg[0]})
+        pta.add_transition('IDLE', 'TX', 'send', energy=3, duration=10,
+                           energy_function=lambda param, arg: 3 + 5 * arg[1],
+                           duration_function=lambda param, arg: 48 + 8 * arg[1])
+        pta.add_transition('TX', 'IDLE', 'txComplete', timeout=2000, is_interrupt=True)
         trace = [
             ['init'],
             ['setTxPower', 10],
@@ -691,24 +690,23 @@ class TestPTA(unittest.TestCase):
         self.assertEqual(duration, expected_duration)
         self.assertEqual(state.name, 'IDLE')
         self.assertEqual(parameters, {
-            'txpower' : 10,
-            'length' : None
+            'txpower': 10,
+            'length': None
         })
 
-
     def test_simulation_param_function(self):
-        pta = PTA(parameters = ['length', 'txpower'])
-        pta.add_state('IDLE', power = 5)
-        pta.add_state('TX', power = 100,
-            power_function = lambda param, arg: 1000 + 2 * param[1])
-        pta.add_transition('UNINITIALIZED', 'IDLE', 'init', energy = 500000, duration = 50000)
-        pta.add_transition('IDLE', 'IDLE', 'setTxPower', energy = 10000, duration = 120,
-            param_update_function = lambda param, arg: {**param, 'txpower' : arg[0]})
-        pta.add_transition('IDLE', 'TX', 'send', energy = 3, duration = 10,
-            energy_function = lambda param, arg: 3 + 5 * arg[1],
-            param_update_function = lambda param, arg: {**param, 'length' : arg[1]})
-        pta.add_transition('TX', 'IDLE', 'txComplete', timeout = 2000, is_interrupt = True,
-            timeout_function = lambda param, arg: 500 + 16 * param[0])
+        pta = PTA(parameters=['length', 'txpower'])
+        pta.add_state('IDLE', power=5)
+        pta.add_state('TX', power=100,
+                      power_function=lambda param, arg: 1000 + 2 * param[1])
+        pta.add_transition('UNINITIALIZED', 'IDLE', 'init', energy=500000, duration=50000)
+        pta.add_transition('IDLE', 'IDLE', 'setTxPower', energy=10000, duration=120,
+                           param_update_function=lambda param, arg: {**param, 'txpower': arg[0]})
+        pta.add_transition('IDLE', 'TX', 'send', energy=3, duration=10,
+                           energy_function=lambda param, arg: 3 + 5 * arg[1],
+                           param_update_function=lambda param, arg: {**param, 'length': arg[1]})
+        pta.add_transition('TX', 'IDLE', 'txComplete', timeout=2000, is_interrupt=True,
+                           timeout_function=lambda param, arg: 500 + 16 * param[0])
         trace = [
             ['init'],
             ['setTxPower', 10],
@@ -721,19 +719,18 @@ class TestPTA(unittest.TestCase):
         self.assertEqual(duration, expected_duration)
         self.assertEqual(state.name, 'IDLE')
         self.assertEqual(parameters, {
-            'txpower' : 10,
-            'length' : 3
+            'txpower': 10,
+            'length': 3
         })
 
     def test_get_X_expensive_state(self):
         pta = PTA.from_json(example_json_1)
         self.assertEqual(pta.get_least_expensive_state(), pta.state['IDLE'])
         self.assertEqual(pta.get_most_expensive_state(), pta.state['TX'])
-        self.assertAlmostEqual(pta.min_duration_until_energy_overflow(), (2**32-1) * 1e-12 / 10e-3, places=9)
-        self.assertAlmostEqual(pta.min_duration_until_energy_overflow(energy_granularity = 1e-9), (2**32-1) * 1e-9 / 10e-3, places=9)
-        self.assertAlmostEqual(pta.max_duration_until_energy_overflow(), (2**32-1) * 1e-12 / 5e-6, places=9)
-        self.assertAlmostEqual(pta.max_duration_until_energy_overflow(energy_granularity = 1e-9), (2**32-1) * 1e-9 / 5e-6, places=9)
-
+        self.assertAlmostEqual(pta.min_duration_until_energy_overflow(), (2**32 - 1) * 1e-12 / 10e-3, places=9)
+        self.assertAlmostEqual(pta.min_duration_until_energy_overflow(energy_granularity=1e-9), (2**32 - 1) * 1e-9 / 10e-3, places=9)
+        self.assertAlmostEqual(pta.max_duration_until_energy_overflow(), (2**32 - 1) * 1e-12 / 5e-6, places=9)
+        self.assertAlmostEqual(pta.max_duration_until_energy_overflow(energy_granularity=1e-9), (2**32 - 1) * 1e-9 / 5e-6, places=9)
 
 
 if __name__ == '__main__':