Added initialize_and_wait

Author: Maxusmusti

docs/index.html

@@ -439,6 +439,33 @@ <h1 class="title">Module <code>state_signals</code></h1>
         self.redis.publish(channel=&#34;event-signal-pubsub&#34;, message=sig.to_json_str())
         self.logger.debug(&#34;Initialization successful!&#34;)
 
+    def initialize_and_wait(
+        self,
+        await_sub_count: int,
+        legal_events: List[str],
+        tag: str = None,
+        expected_resps: List[str] = None,
+        timeout: int = 60,
+    ) -&gt; int:
+        &#34;&#34;&#34;
+        Calls the SignalExporter&#39;s initialize() method, and awaits a specified
+        number of subscribers. Also includes an optional timeout. Returns 0 if
+        sub(s) received, 1 if timed-out (or hangs if no timeout).
+        &#34;&#34;&#34;
+        self.initialize(
+            legal_events=legal_events, tag=tag, expected_resps=expected_resps
+        )
+        counter = 0
+        while len(self.subs) &lt; await_sub_count:
+            time.sleep(0.1)
+            counter += 1
+            if counter &gt;= timeout * 10:
+                self.logger.error(
+                    f&#34;Timeout after waiting {timeout} seconds for {await_sub_count }subs, got {len(self.subs)}&#34;
+                )
+                return 1
+        return 0
+
     def shutdown(self, tag: str = None) -&gt; None:
         &#34;&#34;&#34;
         Shuts down initialization response listener (stops accepting subscribers).
@@ -1135,6 +1162,33 @@ <h3>Methods</h3>
         self.redis.publish(channel=&#34;event-signal-pubsub&#34;, message=sig.to_json_str())
         self.logger.debug(&#34;Initialization successful!&#34;)
 
+    def initialize_and_wait(
+        self,
+        await_sub_count: int,
+        legal_events: List[str],
+        tag: str = None,
+        expected_resps: List[str] = None,
+        timeout: int = 60,
+    ) -&gt; int:
+        &#34;&#34;&#34;
+        Calls the SignalExporter&#39;s initialize() method, and awaits a specified
+        number of subscribers. Also includes an optional timeout. Returns 0 if
+        sub(s) received, 1 if timed-out (or hangs if no timeout).
+        &#34;&#34;&#34;
+        self.initialize(
+            legal_events=legal_events, tag=tag, expected_resps=expected_resps
+        )
+        counter = 0
+        while len(self.subs) &lt; await_sub_count:
+            time.sleep(0.1)
+            counter += 1
+            if counter &gt;= timeout * 10:
+                self.logger.error(
+                    f&#34;Timeout after waiting {timeout} seconds for {await_sub_count }subs, got {len(self.subs)}&#34;
+                )
+                return 1
+        return 0
+
     def shutdown(self, tag: str = None) -&gt; None:
         &#34;&#34;&#34;
         Shuts down initialization response listener (stops accepting subscribers).
@@ -1187,6 +1241,45 @@ <h3>Methods</h3>
     self.logger.debug(&#34;Initialization successful!&#34;)</code></pre>
 </details>
 </dd>
+<dt id="state_signals.SignalExporter.initialize_and_wait"><code class="name flex">
+<span>def <span class="ident">initialize_and_wait</span></span>(<span>self, await_sub_count: int, legal_events: List[str], tag: str = None, expected_resps: List[str] = None, timeout: int = 60) ‑> int</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Calls the SignalExporter's initialize() method, and awaits a specified
+number of subscribers. Also includes an optional timeout. Returns 0 if
+sub(s) received, 1 if timed-out (or hangs if no timeout).</p></div>
+<details class="source">
+<summary>
+<span>Expand source code</span>
+</summary>
+<pre><code class="python">def initialize_and_wait(
+    self,
+    await_sub_count: int,
+    legal_events: List[str],
+    tag: str = None,
+    expected_resps: List[str] = None,
+    timeout: int = 60,
+) -&gt; int:
+    &#34;&#34;&#34;
+    Calls the SignalExporter&#39;s initialize() method, and awaits a specified
+    number of subscribers. Also includes an optional timeout. Returns 0 if
+    sub(s) received, 1 if timed-out (or hangs if no timeout).
+    &#34;&#34;&#34;
+    self.initialize(
+        legal_events=legal_events, tag=tag, expected_resps=expected_resps
+    )
+    counter = 0
+    while len(self.subs) &lt; await_sub_count:
+        time.sleep(0.1)
+        counter += 1
+        if counter &gt;= timeout * 10:
+            self.logger.error(
+                f&#34;Timeout after waiting {timeout} seconds for {await_sub_count }subs, got {len(self.subs)}&#34;
+            )
+            return 1
+    return 0</code></pre>
+</details>
+</dd>
 <dt id="state_signals.SignalExporter.publish_signal"><code class="name flex">
 <span>def <span class="ident">publish_signal</span></span>(<span>self, event: str, sample: int = -1, tag: str = None, metadata: Dict = None, timeout: int = 20) ‑> int</span>
 </code></dt>
@@ -1622,6 +1715,7 @@ <h4><code><a title="state_signals.Signal" href="#state_signals.Signal">Signal</a
 <h4><code><a title="state_signals.SignalExporter" href="#state_signals.SignalExporter">SignalExporter</a></code></h4>
 <ul class="">
 <li><code><a title="state_signals.SignalExporter.initialize" href="#state_signals.SignalExporter.initialize">initialize</a></code></li>
+<li><code><a title="state_signals.SignalExporter.initialize_and_wait" href="#state_signals.SignalExporter.initialize_and_wait">initialize_and_wait</a></code></li>
 <li><code><a title="state_signals.SignalExporter.publish_signal" href="#state_signals.SignalExporter.publish_signal">publish_signal</a></code></li>
 <li><code><a title="state_signals.SignalExporter.shutdown" href="#state_signals.SignalExporter.shutdown">shutdown</a></code></li>
 </ul>

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "state-signals"
-version = "0.2.2"
+version = "0.3.0"
 description = "Package for easy management of state/event signal publishing, subscribing, and responding"
 
 license = "GPL-3.0-only"

state_signals.py

@@ -391,6 +391,33 @@ def initialize(
         self.redis.publish(channel="event-signal-pubsub", message=sig.to_json_str())
         self.logger.debug("Initialization successful!")
 
+    def initialize_and_wait(
+        self,
+        await_sub_count: int,
+        legal_events: List[str],
+        tag: str = None,
+        expected_resps: List[str] = None,
+        timeout: int = 60,
+    ) -> int:
+        """
+        Calls the SignalExporter's initialize() method, and awaits a specified
+        number of subscribers. Also includes an optional timeout. Returns 0 if
+        sub(s) received, 1 if timed-out (or hangs if no timeout).
+        """
+        self.initialize(
+            legal_events=legal_events, tag=tag, expected_resps=expected_resps
+        )
+        counter = 0
+        while len(self.subs) < await_sub_count:
+            time.sleep(0.1)
+            counter += 1
+            if counter >= timeout * 10:
+                self.logger.error(
+                    f"Timeout after waiting {timeout} seconds for {await_sub_count }subs, got {len(self.subs)}"
+                )
+                return 1
+        return 0
+
     def shutdown(self, tag: str = None) -> None:
         """
         Shuts down initialization response listener (stops accepting subscribers).

tests/basic_test.py

@@ -31,7 +31,9 @@ def _listener(responder):
 
 def _init():
     resp_proc.start()
-    sig_ex.initialize(legal_events=["benchmark-start", "benchmark-stop"])
+    return sig_ex.initialize_and_wait(
+        1, legal_events=["benchmark-start", "benchmark-stop"]
+    )
 
 
 def _shutdown():
@@ -44,9 +46,9 @@ def _cleanup():
 
 @pytest.mark.dependency()
 def test_init():
-    _init()
+    sub_check = _init()
     assert sig_ex.init_listener.is_alive()
-    time.sleep(1)
+    assert sub_check == 0
     assert sig_ex.subs