From c9de1ee757d7943bbd51231a7eec8f3e45b69467 Mon Sep 17 00:00:00 2001
From: meejah <meejah@meejah.ca>
Date: Tue, 12 May 2020 14:22:32 -0600
Subject: [PATCH] add a .restart() to Client and StorageServer

---
 integration/grid.py              | 20 ++++++++++++++++++--
 integration/test_grid_manager.py | 19 ++++---------------
 2 files changed, 22 insertions(+), 17 deletions(-)

diff --git a/integration/grid.py b/integration/grid.py
index 166cf940d..d7b104e16 100644
--- a/integration/grid.py
+++ b/integration/grid.py
@@ -48,6 +48,7 @@ from util import (
     _cleanup_tahoe_process,
     _tahoe_runner_optional_coverage,
     TahoeProcess,
+    await_client_ready,
 )
 
 import attr
@@ -155,8 +156,14 @@ class StorageServer(object):
         validator=attr.validators.provides(IProcessProtocol)
     )
 
-    # XXX needs a restart() probably .. or at least a stop() and
-    # start()
+    @inlineCallbacks
+    def restart(self, reactor, request):
+        self.process.transport.signalProcess('TERM')
+        yield self.protocol.exited
+        self.process = yield _run_node(
+            reactor, self.process.node_dir, request, None,
+        )
+        self.protocol = self.process.transport._protocol
 
 
 @inlineCallbacks
@@ -190,6 +197,14 @@ class Client(object):
         validator=attr.validators.provides(IProcessProtocol)
     )
 
+    @inlineCallbacks
+    def restart(self, reactor, request):
+        self.process.transport.signalProcess('TERM')
+        yield self.protocol.exited
+        x = yield _run_node(
+            reactor, self.process.node_dir, request, None,
+        )
+
     # XXX add stop / start / restart
     # ...maybe "reconfig" of some kind?
 
@@ -373,6 +388,7 @@ class Grid(object):
             total=total,
         )
         self.clients[name] = client
+        yield await_client_ready(client.process)
         returnValue(client)
 
 
diff --git a/integration/test_grid_manager.py b/integration/test_grid_manager.py
index 863a806a6..38c2dd78a 100644
--- a/integration/test_grid_manager.py
+++ b/integration/test_grid_manager.py
@@ -133,8 +133,6 @@ def test_reject_storage_server(reactor, request, temp_dir, flog_gatherer, port_a
     )
     assert json.loads(gm_config)['storage_servers'].keys() == ['storage0']
 
-    diana = yield grid.add_client("diana", needed=2, happy=2, total=2)
-
     print("inserting certificate")
     cert = yield util.run_tahoe(
         reactor, request, "grid-manager", "--config", "-", "sign", "storage0",
@@ -150,29 +148,20 @@ def test_reject_storage_server(reactor, request, temp_dir, flog_gatherer, port_a
         config.write(f)
 
     # re-start this storage server
-    storage0.process.transport.signalProcess('TERM')
-    yield storage0.protocol.exited
-    yield util._run_node(
-        reactor, storage0.process.node_dir, request, None,
-    )
-
-    yield util.await_client_ready(diana.process, servers=2)
+    yield storage0.restart(reactor, request)
 
     # now only one storage-server has the certificate .. configure
     # diana to have the grid-manager certificate
 
+    diana = yield grid.add_client("diana", needed=2, happy=2, total=2)
+
     config = configutil.get_config(join(diana.process.node_dir, "tahoe.cfg"))
     config.add_section("grid_managers")
     config.set("grid_managers", "test", ed25519.string_from_verifying_key(gm_pubkey))
     with open(join(diana.process.node_dir, "tahoe.cfg"), "w") as f:
         config.write(f)
-    diana.process.transport.signalProcess('TERM')
-    yield diana.protocol.exited
 
-    diana = yield util._run_node(
-        reactor, diana.process.node_dir, request, None,
-    )
-    yield util.await_client_ready(diana.process, servers=2)
+    yield diana.restart(reactor, request)
 
     # try to put something into the grid, which should fail (because
     # diana has happy=2 but should only find storage0 to be acceptable