From c999953ccb21054fd7ac5d00dd30e3b65b7255ce Mon Sep 17 00:00:00 2001
From: Christopher Amin <camin@ripe.net>
Date: Fri, 18 Nov 2022 15:40:08 +0100
Subject: [PATCH] 3.0.3: fix stream waiting bugs

---
 CHANGES.rst                   |  4 ++++
 ripe/atlas/tools/streaming.py | 11 ++++++-----
 ripe/atlas/tools/version.py   |  2 +-
 3 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/CHANGES.rst b/CHANGES.rst
index 75078b3..c0d1133 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -1,5 +1,9 @@
 Release History
 ===============
+3.0.3 (release 2022-11-18)
+--------------------------
+- Fix issue where the measure command would continue to stream results after all probes have responded
+
 3.0.2 (release 2022-05-23)
 --------------------------
 - Fix "measure spec" command which was broken due to cousteau issue
diff --git a/ripe/atlas/tools/streaming.py b/ripe/atlas/tools/streaming.py
index 7182ead..bd7e20f 100644
--- a/ripe/atlas/tools/streaming.py
+++ b/ripe/atlas/tools/streaming.py
@@ -19,7 +19,7 @@
 from .settings import conf
 
 
-class CaptureLimitExceeded(Exception):
+class CaptureLimitMet(Exception):
     pass
 
 
@@ -62,11 +62,12 @@ def on_result_response(result, *args):
                 stream.timeout(self.STREAM_INTERVAL)
                 for result in results:
                     self.num_results += 1
-                    if self.capture_limit and self.num_results > self.capture_limit:
-                        raise CaptureLimitExceeded()
                     yield result
+                    if self.capture_limit and self.num_results == self.capture_limit:
+                        raise CaptureLimitMet()
                 results = []
                 if self.timeout is not None:
                     remaining = start + self.timeout - time.time()
-        except (KeyboardInterrupt, CaptureLimitExceeded):
-            stream.disconnect()
+        except (KeyboardInterrupt, CaptureLimitMet):
+            pass
+        stream.disconnect()
diff --git a/ripe/atlas/tools/version.py b/ripe/atlas/tools/version.py
index 10e290f..2c84278 100644
--- a/ripe/atlas/tools/version.py
+++ b/ripe/atlas/tools/version.py
@@ -13,4 +13,4 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
-__version__ = "3.0.2"
+__version__ = "3.0.3"