1  """
  2  Tests for client connect-point functionality.
  3  """
  4  
  5  from arti_rpc_tests import arti_test
  6  from arti_rpc import (
  7      ArtiRpcError,
  8      ArtiRpcErrorStatus,
  9      ArtiRpcConn,
 10      ArtiRpcConnBuilder,
 11  )
 12  
 13  import tempfile
 14  import os
 15  import urllib.parse
 16  
 17  
 18  def connpt_abort():
 19      """
 20      Return a connect point that fails with an explicit abort.
 21      """
 22      return """\
 23  [builtin]
 24  builtin = "abort"
 25  """
 26  
 27  
 28  def connpt_unix(context):
 29      """
 30      Return a connect point that uses an AF_UNIX connection
 31      to connect to arti.
 32      """
 33      return f"""
 34  [connect]
 35  socket = "unix:{context.socket_path}"
 36  auth = "none"
 37  """
 38  
 39  
 40  def connpt_working(context):
 41      """
 42      Return a connect point that should work for connecting to the
 43      arti instance in `context`.
 44  
 45      (Prefer this function to `connpt_unix()`, so that once we have
 46      something that works on windows, we can make our tests pass there too.)
 47      """
 48      return connpt_unix(context)
 49  
 50  
 51  class Tempdir:
 52      """
 53      Helper for creating files within a temporary directory.
 54  
 55      When this object is destroyed, the directory and its contents are deleted.
 56      """
 57  
 58      def __init__(self):
 59          self.tmpdir = tempfile.TemporaryDirectory()
 60  
 61      def dirname(self):
 62          """
 63          Return the name of the temporary directory
 64          """
 65          return self.tmpdir.name
 66  
 67      def fname(self, name):
 68          """
 69          Return the name of a file within the temporary directory
 70          """
 71          return os.path.join(self.tmpdir.name, name)
 72  
 73      def write(self, fname, text):
 74          """
 75          Store `text` into the file called `fname` within the temporary directory.
 76          """
 77          with open(self.fname(fname), "w") as f:
 78              f.write(text)
 79  
 80  
 81  class SavedEnviron:
 82      """
 83      Context manager to preserve os.environ while a test is changing it.
 84      """
 85  
 86      def __init__(self):
 87          pass
 88  
 89      def __enter__(self):
 90          self.env = dict(os.environ)
 91  
 92      def __exit__(self, _et, _ev, _tb):
 93          os.environ.update(self.env)
 94          for k in os.environ.keys() - self.env.keys():
 95              del os.environ[k]
 96  
 97  
 98  def assert_builder_aborts(bld: ArtiRpcConnBuilder):
 99      """
100      Try to open a RPC connection with `bld`, and assert that the attempt aborts.
101      """
102      try:
103          _ = bld.connect()
104          assert False  # shouldn't be reached.
105      except ArtiRpcError as e:
106          assert e.status_code() == ArtiRpcErrorStatus.ALL_CONNECT_ATTEMPTS_FAILED
107          assert 'Encountered an explicit "abort"' in str(e)
108  
109  
110  def assert_builder_connects(bld):
111      """
112      Try to open a RPC connection with `bld`, and assert that the attempt succeeds.
113      """
114      c = bld.connect()
115      assert c is not None
116  
117  
118  @arti_test
119  def ordering_literal_manual(context):
120      # Assert that prepend_literal_connect_point respects ordering.
121      bld = ArtiRpcConnBuilder()
122      bld.prepend_literal_connect_point(connpt_working(context))
123      bld.prepend_literal_connect_point(connpt_abort())
124      assert_builder_aborts(bld)
125  
126      bld = ArtiRpcConnBuilder()
127      bld.prepend_literal_connect_point(connpt_abort())
128      bld.prepend_literal_connect_point(connpt_working(context))
129      assert_builder_connects(bld)
130  
131  
132  @arti_test
133  def ordering_paths_manual(context):
134      # Assert that prepend_literal_path respects ordering.
135      tmp = Tempdir()
136      tmp.write("abort.toml", connpt_abort())
137      tmp.write("working.toml", connpt_working(context))
138  
139      bld = ArtiRpcConnBuilder()
140      bld.prepend_literal_path(tmp.fname("working.toml"))
141      bld.prepend_literal_path(tmp.fname("abort.toml"))
142      assert_builder_aborts(bld)
143  
144      bld = ArtiRpcConnBuilder()
145      bld.prepend_literal_path(tmp.fname("abort.toml"))
146      bld.prepend_literal_path(tmp.fname("working.toml"))
147      assert_builder_connects(bld)
148  
149  
150  @arti_test
151  def ordering_env_path(context):
152      # Assert that paths within our envvars respect ordering.
153      tmp = Tempdir()
154      tmp.write("abort.toml", connpt_abort())
155      tmp.write("working.toml", connpt_working(context))
156  
157      fn_a = tmp.fname("abort.toml")
158      fn_w = tmp.fname("working.toml")
159      assert ":" not in fn_a
160      assert ":" not in fn_w
161  
162      for varname in ["ARTI_RPC_CONNECT_PATH", "ARTI_RPC_CONNECT_PATH_OVERRIDE"]:
163          assert varname not in os.environ
164          with SavedEnviron():
165              bld = ArtiRpcConnBuilder()
166              os.environ[varname] = f"{fn_a}:{fn_w}"
167              assert_builder_aborts(bld)
168  
169              bld = ArtiRpcConnBuilder()
170              os.environ[varname] = f"{fn_w}:{fn_a}"
171              assert_builder_connects(bld)
172          assert varname not in os.environ
173  
174  
175  @arti_test
176  def ordering_env_literal(context):
177      # Assert that literal connect points within our envvars respect ordering.
178      q_a = urllib.parse.quote(connpt_abort())
179      q_w = urllib.parse.quote(connpt_working(context))
180  
181      for varname in ["ARTI_RPC_CONNECT_PATH", "ARTI_RPC_CONNECT_PATH_OVERRIDE"]:
182          assert varname not in os.environ
183          with SavedEnviron():
184              bld = ArtiRpcConnBuilder()
185              os.environ[varname] = f"{q_a}:{q_w}"
186              assert_builder_aborts(bld)
187  
188              bld = ArtiRpcConnBuilder()
189              os.environ[varname] = f"{q_w}:{q_a}"
190              assert_builder_connects(bld)
191          assert varname not in os.environ
192  
193  
194  @arti_test
195  def ordering_dir(context):
196      # Assert that files within a directory respect ordering.
197      tmp = Tempdir()
198      tmp.write("00_name_ignored", connpt_abort())  # no ".toml" suffix
199      tmp.write("01_abort.toml", connpt_abort())
200      tmp.write("02_connect.toml", connpt_working(context))
201      tmp.write("03_abort.toml", connpt_abort())
202  
203      bld = ArtiRpcConnBuilder()
204      bld.prepend_literal_path(tmp.dirname())
205      assert_builder_aborts(bld)
206  
207      os.unlink(tmp.fname("01_abort.toml"))
208      assert_builder_connects(bld)
209  
210  
211  @arti_test
212  def ordering_multi(context):
213      # Assert that our envvars and manually inserted items respect ordering
214      # wrt one another.
215      tmp = Tempdir()
216      tmp.write("abort.toml", connpt_abort())
217      tmp.write("working.toml", connpt_working(context))
218  
219      fn_a = tmp.fname("abort.toml")
220      fn_w = tmp.fname("working.toml")
221  
222      with SavedEnviron():
223          bld = ArtiRpcConnBuilder()
224          os.environ["ARTI_RPC_CONNECT_PATH"] = fn_a
225          assert_builder_aborts(bld)
226  
227          bld.prepend_literal_path(fn_w)
228          assert_builder_connects(bld)
229  
230          os.environ["ARTI_RPC_CONNECT_PATH_OVERRIDE"] = fn_a
231          assert_builder_aborts(bld)
232  
233  
234  @arti_test
235  def connect_nobuilder(context):
236      # Assert that we can use ArtiRpcConn constructor
237      # without a builder.
238      tmp = Tempdir()
239      tmp.write("abort.toml", connpt_abort())
240      tmp.write("working.toml", connpt_working(context))
241  
242      fn_a = tmp.fname("abort.toml")
243      fn_w = tmp.fname("working.toml")
244  
245      with SavedEnviron():
246          os.environ["ARTI_RPC_CONNECT_PATH"] = fn_a
247          try:
248              _ = ArtiRpcConn()
249              assert False  # Shouldn't be reached.
250          except ArtiRpcError as e:
251              assert e.status_code() == ArtiRpcErrorStatus.ALL_CONNECT_ATTEMPTS_FAILED
252              assert 'Encountered an explicit "abort"' in str(e)
253  
254          os.environ["ARTI_RPC_CONNECT_PATH"] = fn_w
255          c = ArtiRpcConn()
256          assert c is not None