Commit 7d68444d authored by Robin Obůrka's avatar Robin Obůrka Committed by Robin Obůrka

msgloop: Overwrite teardown_box too

parent eb837beb
...@@ -147,6 +147,11 @@ class SNPipelineBox(SNBox): ...@@ -147,6 +147,11 @@ class SNPipelineBox(SNBox):
if not self.socket_recv and not self.socket_send: if not self.socket_recv and not self.socket_send:
raise SetupError("Neither input nor output socket provided") raise SetupError("Neither input nor output socket provided")
def teardown_box(self):
self.socket_recv.close()
self.socket_send.close()
super().teardown_box()
def get_processed_message(self): def get_processed_message(self):
msg = self.socket_recv.recv_multipart() msg = self.socket_recv.recv_multipart()
msg_type, payload = parse_msg(msg) msg_type, payload = parse_msg(msg)
...@@ -183,6 +188,10 @@ class SNGeneratorBox(SNBox): ...@@ -183,6 +188,10 @@ class SNGeneratorBox(SNBox):
if not inspect.isgeneratorfunction(self.process): if not inspect.isgeneratorfunction(self.process):
raise SetupError("Generator is expected for output-only box") raise SetupError("Generator is expected for output-only box")
def teardown_box(self):
self.socket_send.close()
super().teardown_box()
def get_processed_message(self): def get_processed_message(self):
return next(self.process_iterator) return next(self.process_iterator)
...@@ -209,6 +218,10 @@ class SNTerminationBox(SNBox): ...@@ -209,6 +218,10 @@ class SNTerminationBox(SNBox):
if not self.socket_recv: if not self.socket_recv:
raise SetupError("Input socket wasn't provided") raise SetupError("Input socket wasn't provided")
def teardown_box(self):
self.socket_recv.close()
super().teardown_box()
def get_processed_message(self): def get_processed_message(self):
msg = self.socket_recv.recv_multipart() msg = self.socket_recv.recv_multipart()
msg_type, payload = parse_msg(msg) msg_type, payload = parse_msg(msg)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment