prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  mock_xray_client cannot be used as a context manager
  PR #1255 added support for `aws_xray_sdk` which is great.
  But there is a problem with it: `moto.mock_xray_client` is *only* a function decorator, and unlike all other `mock_*` methods it cannot be used as a context manager or directly with `start()`...`stop()`.
  As a result, it is not possible to write a `py.test` fixture which would add support for mocking `xray_client`.
  
  Also, `mock_xray_client` does not return the result of the function it decorates. Given it is meant to be used to decorate test functions it is most likely not a big issue, but I think it is still worth fixing.
  
  I will prepare a PR for the return value issue soon.
  Also I am thinking about refactoring `mock_xray_client` to base it on the existing infrastructure (`BaseBackend`, `base_decorator`) but am not yet enough familiar with `moto` internals to be sure which would be the best way to implement it.
  
  Installed version: `moto-ext==1.1.25`
  The problem seemingly persists in current `master` branch.
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `moto/xray/mock_client.py`
  Content:
  ```
  1 from functools import wraps
  2 import os
  3 from moto.xray import xray_backends
  4 import aws_xray_sdk.core
  5 from aws_xray_sdk.core.context import Context as AWSContext
  6 from aws_xray_sdk.core.emitters.udp_emitter import UDPEmitter
  7 
  8 
  9 class MockEmitter(UDPEmitter):
  10     """
  11     Replaces the code that sends UDP to local X-Ray daemon
  12     """
  13     def __init__(self, daemon_address='127.0.0.1:2000'):
  14         address = os.getenv('AWS_XRAY_DAEMON_ADDRESS_YEAH_NOT_TODAY_MATE', daemon_address)
  15         self._ip, self._port = self._parse_address(address)
  16 
  17     def _xray_backend(self, region):
  18         return xray_backends[region]
  19 
  20     def send_entity(self, entity):
  21         # Hack to get region
  22         # region = entity.subsegments[0].aws['region']
  23         # xray = self._xray_backend(region)
  24 
  25         # TODO store X-Ray data, pretty sure X-Ray needs refactor for this
  26         pass
  27 
  28     def _send_data(self, data):
  29         raise RuntimeError('Should not be running this')
  30 
  31 
  32 def mock_xray_client(f):
  33     """
  34     Mocks the X-Ray sdk by pwning its evil singleton with our methods
  35 
  36     The X-Ray SDK has normally been imported and `patched()` called long before we start mocking.
  37     This means the Context() will be very unhappy if an env var isnt present, so we set that, save
  38     the old context, then supply our new context.
  39     We also patch the Emitter by subclassing the UDPEmitter class replacing its methods and pushing
  40     that itno the recorder instance.
  41     """
  42     @wraps(f)
  43     def _wrapped(*args, **kwargs):
  44         print("Starting X-Ray Patch")
  45 
  46         old_xray_context_var = os.environ.get('AWS_XRAY_CONTEXT_MISSING')
  47         os.environ['AWS_XRAY_CONTEXT_MISSING'] = 'LOG_ERROR'
  48         old_xray_context = aws_xray_sdk.core.xray_recorder._context
  49         old_xray_emitter = aws_xray_sdk.core.xray_recorder._emitter
  50         aws_xray_sdk.core.xray_recorder._context = AWSContext()
  51         aws_xray_sdk.core.xray_recorder._emitter = MockEmitter()
  52 
  53         try:
  54             f(*args, **kwargs)
  55         finally:
  56 
  57             if old_xray_context_var is None:
  58                 del os.environ['AWS_XRAY_CONTEXT_MISSING']
  59             else:
  60                 os.environ['AWS_XRAY_CONTEXT_MISSING'] = old_xray_context_var
  61 
  62             aws_xray_sdk.core.xray_recorder._emitter = old_xray_emitter
  63             aws_xray_sdk.core.xray_recorder._context = old_xray_context
  64 
  65     return _wrapped
  66 
  67 
  68 class XRaySegment(object):
  69     """
  70     XRay is request oriented, when a request comes in, normally middleware like django (or automatically in lambda) will mark
  71     the start of a segment, this stay open during the lifetime of the request. During that time subsegments may be generated
  72     by calling other SDK aware services or using some boto functions. Once the request is finished, middleware will also stop
  73     the segment, thus causing it to be emitted via UDP.
  74 
  75     During testing we're going to have to control the start and end of a segment via context managers.
  76     """
  77     def __enter__(self):
  78         aws_xray_sdk.core.xray_recorder.begin_segment(name='moto_mock', traceid=None, parent_id=None, sampling=1)
  79 
  80         return self
  81 
  82     def __exit__(self, exc_type, exc_val, exc_tb):
  83         aws_xray_sdk.core.xray_recorder.end_segment()
  84
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: diff --git a/moto/xray/mock_client.py b/moto/xray/mock_client.py
  --- a/moto/xray/mock_client.py
  +++ b/moto/xray/mock_client.py
  @@ -51,7 +51,7 @@
           aws_xray_sdk.core.xray_recorder._emitter = MockEmitter()
   
           try:
  -            f(*args, **kwargs)
  +            return f(*args, **kwargs)
           finally:
   
               if old_xray_context_var is None:
  
verification_info: {"golden_diff": "diff --git a/moto/xray/mock_client.py b/moto/xray/mock_client.py\n--- a/moto/xray/mock_client.py\n+++ b/moto/xray/mock_client.py\n@@ -51,7 +51,7 @@\n         aws_xray_sdk.core.xray_recorder._emitter = MockEmitter()\n \n         try:\n-            f(*args, **kwargs)\n+            return f(*args, **kwargs)\n         finally:\n \n             if old_xray_context_var is None:\n", "issue": "mock_xray_client cannot be used as a context manager\nPR #1255 added support for `aws_xray_sdk` which is great.\r\nBut there is a problem with it: `moto.mock_xray_client` is *only* a function decorator, and unlike all other `mock_*` methods it cannot be used as a context manager or directly with `start()`...`stop()`.\r\nAs a result, it is not possible to write a `py.test` fixture which would add support for mocking `xray_client`.\r\n\r\nAlso, `mock_xray_client` does not return the result of the function it decorates. Given it is meant to be used to decorate test functions it is most likely not a big issue, but I think it is still worth fixing.\r\n\r\nI will prepare a PR for the return value issue soon.\r\nAlso I am thinking about refactoring `mock_xray_client` to base it on the existing infrastructure (`BaseBackend`, `base_decorator`) but am not yet enough familiar with `moto` internals to be sure which would be the best way to implement it.\r\n\r\nInstalled version: `moto-ext==1.1.25`\r\nThe problem seemingly persists in current `master` branch.\n", "before_files": [{"content": "from functools import wraps\nimport os\nfrom moto.xray import xray_backends\nimport aws_xray_sdk.core\nfrom aws_xray_sdk.core.context import Context as AWSContext\nfrom aws_xray_sdk.core.emitters.udp_emitter import UDPEmitter\n\n\nclass MockEmitter(UDPEmitter):\n    \"\"\"\n    Replaces the code that sends UDP to local X-Ray daemon\n    \"\"\"\n    def __init__(self, daemon_address='127.0.0.1:2000'):\n        address = os.getenv('AWS_XRAY_DAEMON_ADDRESS_YEAH_NOT_TODAY_MATE', daemon_address)\n        self._ip, self._port = self._parse_address(address)\n\n    def _xray_backend(self, region):\n        return xray_backends[region]\n\n    def send_entity(self, entity):\n        # Hack to get region\n        # region = entity.subsegments[0].aws['region']\n        # xray = self._xray_backend(region)\n\n        # TODO store X-Ray data, pretty sure X-Ray needs refactor for this\n        pass\n\n    def _send_data(self, data):\n        raise RuntimeError('Should not be running this')\n\n\ndef mock_xray_client(f):\n    \"\"\"\n    Mocks the X-Ray sdk by pwning its evil singleton with our methods\n\n    The X-Ray SDK has normally been imported and `patched()` called long before we start mocking.\n    This means the Context() will be very unhappy if an env var isnt present, so we set that, save\n    the old context, then supply our new context.\n    We also patch the Emitter by subclassing the UDPEmitter class replacing its methods and pushing\n    that itno the recorder instance.\n    \"\"\"\n    @wraps(f)\n    def _wrapped(*args, **kwargs):\n        print(\"Starting X-Ray Patch\")\n\n        old_xray_context_var = os.environ.get('AWS_XRAY_CONTEXT_MISSING')\n        os.environ['AWS_XRAY_CONTEXT_MISSING'] = 'LOG_ERROR'\n        old_xray_context = aws_xray_sdk.core.xray_recorder._context\n        old_xray_emitter = aws_xray_sdk.core.xray_recorder._emitter\n        aws_xray_sdk.core.xray_recorder._context = AWSContext()\n        aws_xray_sdk.core.xray_recorder._emitter = MockEmitter()\n\n        try:\n            f(*args, **kwargs)\n        finally:\n\n            if old_xray_context_var is None:\n                del os.environ['AWS_XRAY_CONTEXT_MISSING']\n            else:\n                os.environ['AWS_XRAY_CONTEXT_MISSING'] = old_xray_context_var\n\n            aws_xray_sdk.core.xray_recorder._emitter = old_xray_emitter\n            aws_xray_sdk.core.xray_recorder._context = old_xray_context\n\n    return _wrapped\n\n\nclass XRaySegment(object):\n    \"\"\"\n    XRay is request oriented, when a request comes in, normally middleware like django (or automatically in lambda) will mark\n    the start of a segment, this stay open during the lifetime of the request. During that time subsegments may be generated\n    by calling other SDK aware services or using some boto functions. Once the request is finished, middleware will also stop\n    the segment, thus causing it to be emitted via UDP.\n\n    During testing we're going to have to control the start and end of a segment via context managers.\n    \"\"\"\n    def __enter__(self):\n        aws_xray_sdk.core.xray_recorder.begin_segment(name='moto_mock', traceid=None, parent_id=None, sampling=1)\n\n        return self\n\n    def __exit__(self, exc_type, exc_val, exc_tb):\n        aws_xray_sdk.core.xray_recorder.end_segment()\n", "path": "moto/xray/mock_client.py"}], "after_files": [{"content": "from functools import wraps\nimport os\nfrom moto.xray import xray_backends\nimport aws_xray_sdk.core\nfrom aws_xray_sdk.core.context import Context as AWSContext\nfrom aws_xray_sdk.core.emitters.udp_emitter import UDPEmitter\n\n\nclass MockEmitter(UDPEmitter):\n    \"\"\"\n    Replaces the code that sends UDP to local X-Ray daemon\n    \"\"\"\n    def __init__(self, daemon_address='127.0.0.1:2000'):\n        address = os.getenv('AWS_XRAY_DAEMON_ADDRESS_YEAH_NOT_TODAY_MATE', daemon_address)\n        self._ip, self._port = self._parse_address(address)\n\n    def _xray_backend(self, region):\n        return xray_backends[region]\n\n    def send_entity(self, entity):\n        # Hack to get region\n        # region = entity.subsegments[0].aws['region']\n        # xray = self._xray_backend(region)\n\n        # TODO store X-Ray data, pretty sure X-Ray needs refactor for this\n        pass\n\n    def _send_data(self, data):\n        raise RuntimeError('Should not be running this')\n\n\ndef mock_xray_client(f):\n    \"\"\"\n    Mocks the X-Ray sdk by pwning its evil singleton with our methods\n\n    The X-Ray SDK has normally been imported and `patched()` called long before we start mocking.\n    This means the Context() will be very unhappy if an env var isnt present, so we set that, save\n    the old context, then supply our new context.\n    We also patch the Emitter by subclassing the UDPEmitter class replacing its methods and pushing\n    that itno the recorder instance.\n    \"\"\"\n    @wraps(f)\n    def _wrapped(*args, **kwargs):\n        print(\"Starting X-Ray Patch\")\n\n        old_xray_context_var = os.environ.get('AWS_XRAY_CONTEXT_MISSING')\n        os.environ['AWS_XRAY_CONTEXT_MISSING'] = 'LOG_ERROR'\n        old_xray_context = aws_xray_sdk.core.xray_recorder._context\n        old_xray_emitter = aws_xray_sdk.core.xray_recorder._emitter\n        aws_xray_sdk.core.xray_recorder._context = AWSContext()\n        aws_xray_sdk.core.xray_recorder._emitter = MockEmitter()\n\n        try:\n            return f(*args, **kwargs)\n        finally:\n\n            if old_xray_context_var is None:\n                del os.environ['AWS_XRAY_CONTEXT_MISSING']\n            else:\n                os.environ['AWS_XRAY_CONTEXT_MISSING'] = old_xray_context_var\n\n            aws_xray_sdk.core.xray_recorder._emitter = old_xray_emitter\n            aws_xray_sdk.core.xray_recorder._context = old_xray_context\n\n    return _wrapped\n\n\nclass XRaySegment(object):\n    \"\"\"\n    XRay is request oriented, when a request comes in, normally middleware like django (or automatically in lambda) will mark\n    the start of a segment, this stay open during the lifetime of the request. During that time subsegments may be generated\n    by calling other SDK aware services or using some boto functions. Once the request is finished, middleware will also stop\n    the segment, thus causing it to be emitted via UDP.\n\n    During testing we're going to have to control the start and end of a segment via context managers.\n    \"\"\"\n    def __enter__(self):\n        aws_xray_sdk.core.xray_recorder.begin_segment(name='moto_mock', traceid=None, parent_id=None, sampling=1)\n\n        return self\n\n    def __exit__(self, exc_type, exc_val, exc_tb):\n        aws_xray_sdk.core.xray_recorder.end_segment()\n", "path": "moto/xray/mock_client.py"}]}