From a0e7438ccc622d909f9b7c1eb367f6f7d2e80c92 Mon Sep 17 00:00:00 2001
From: Demi Marie Obenour <demi@invisiblethingslab.com>
Date: Mon, 22 Apr 2024 21:16:24 -0400
Subject: [PATCH] Make more functions in agent tests idempotent

This will hopefully make it easier to support multiple subtests per test
in the future.
---
 qrexec/tests/socket/agent.py | 15 +++++++--------
 1 file changed, 7 insertions(+), 8 deletions(-)

diff --git a/qrexec/tests/socket/agent.py b/qrexec/tests/socket/agent.py
index 6a8ceab8..eb351719 100644
--- a/qrexec/tests/socket/agent.py
+++ b/qrexec/tests/socket/agent.py
@@ -77,6 +77,7 @@ def make_executable_service(self, *args):
         util.make_executable_service(self.tempdir, *args)
 
     def setUp(self):
+        self.agent = self.dom0 = None
         self.tempdir = tempfile.mkdtemp()
         os.mkdir(os.path.join(self.tempdir, "local-rpc"))
         os.mkdir(os.path.join(self.tempdir, "rpc"))
@@ -84,6 +85,8 @@ def setUp(self):
         self.addCleanup(shutil.rmtree, self.tempdir)
 
     def start_agent(self):
+        if self.agent is not None:
+            return
         env = os.environ.copy()
         env["LD_LIBRARY_PATH"] = os.path.join(ROOT_PATH, "libqrexec")
         env["VCHAN_DOMAIN"] = str(self.domain)
@@ -127,8 +130,11 @@ def wait_for_agent_children(self):
         psutil.wait_procs(children)
 
     def connect_dom0(self):
-        dom0 = qrexec.vchan_client(self.tempdir, self.domain, 0, 512)
+        if self.dom0 is not None:
+            return self.dom0
+        dom0 = self.dom0 = qrexec.vchan_client(self.tempdir, self.domain, 0, 512)
         self.addCleanup(dom0.close)
+        dom0.handshake()
         return dom0
 
     def connect_target(self):
@@ -150,13 +156,11 @@ def test_handshake(self):
         self.start_agent()
 
         dom0 = self.connect_dom0()
-        dom0.handshake()
 
     def _test_just_exec(self, cmd):
         self.start_agent()
 
         dom0 = self.connect_dom0()
-        dom0.handshake()
 
         user = getpass.getuser().encode("ascii")
 
@@ -235,7 +239,6 @@ def test_exec_cmdline(self):
         self.start_agent()
 
         dom0 = self.connect_dom0()
-        dom0.handshake()
 
         user = getpass.getuser().encode("ascii")
 
@@ -260,7 +263,6 @@ def test_trigger_service(self):
         target_domain_name = b"target_domain"
 
         dom0 = self.connect_dom0()
-        dom0.handshake()
 
         client = self.connect_client()
         ident = self.trigger_service(
@@ -286,7 +288,6 @@ def test_trigger_service_refused(self):
         target_domain_name = b"target_domain"
 
         dom0 = self.connect_dom0()
-        dom0.handshake()
 
         client = self.connect_client()
         ident = self.trigger_service(
@@ -337,7 +338,6 @@ def execute_qubesrpc(self, service: str, src_domain_name: str):
         self.start_agent()
 
         dom0 = self.connect_dom0()
-        dom0.handshake()
 
         user = getpass.getuser()
         cmdline = "{}:QUBESRPC {} {}\0".format(
@@ -977,7 +977,6 @@ def execute(self, cmd: str):
         self.start_agent()
 
         dom0 = self.connect_dom0()
-        dom0.handshake()
 
         user = getpass.getuser()
         cmdline = "{}:{}\0".format(user, cmd).encode("ascii")