@@ -114,19 +114,23 @@ def test_src_handler_restricted(self):
114114 self .assertTrue (fsm .states .packets_ready )
115115 self .assertEqual (fsm .states .num_packets_ready , 1 )
116116 next_pdu = self .source_handler .get_next_packet ()
117- self .assertIsInstance (next_pdu .base , MetadataPdu )
117+ assert next_pdu is not None
118+ self .assertIsInstance (next_pdu .pdu , MetadataPdu )
118119 fsm = self .source_handler .state_machine ()
119120 self .assertTrue (fsm .states .packets_ready )
120- file_data = self .source_handler .get_next_packet ()
121- self .assertIsInstance (file_data .base , FileDataPdu )
121+ next_pdu = self .source_handler .get_next_packet ()
122+ assert next_pdu is not None
123+ self .assertIsInstance (next_pdu .pdu , FileDataPdu )
122124 fsm = self .source_handler .state_machine ()
123125 self .assertTrue (fsm .states .packets_ready )
124- eof_data = self .source_handler .get_next_packet ()
125- self .assertIsInstance (eof_data .base , EofPdu )
126+ next_pdu = self .source_handler .get_next_packet ()
127+ assert next_pdu is not None
128+ self .assertIsInstance (next_pdu .pdu , EofPdu )
129+ assert isinstance (next_pdu .pdu , EofPdu )
126130 # Send ACK
127131 pdu_conf = PduConfig (
128132 direction = Direction .TOWARDS_SENDER ,
129- transaction_seq_num = eof_data . base .transaction_seq_num ,
133+ transaction_seq_num = next_pdu . pdu .transaction_seq_num ,
130134 source_entity_id = self .source_id ,
131135 dest_entity_id = self .dest_id ,
132136 trans_mode = TransmissionMode .ACKNOWLEDGED ,
@@ -147,6 +151,7 @@ def test_src_handler_restricted(self):
147151 fsm = self .source_handler .state_machine (packet = finished )
148152 self .assertTrue (fsm .states .packets_ready )
149153 finished_ack = self .source_handler .get_next_packet ()
150- self .assertIsInstance (finished_ack .base , AckPdu )
154+ assert finished_ack is not None
155+ self .assertIsInstance (finished_ack .pdu , AckPdu )
151156 fsm = self .source_handler .state_machine ()
152157 self .assertEqual (fsm .states .state , CfdpState .IDLE )
0 commit comments