|
| 1 | +# Copyright 2010 New Relic, Inc. |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +import functools |
| 16 | +import logging |
| 17 | +import socket |
| 18 | +import threading |
| 19 | + |
| 20 | +import moto.server |
| 21 | +import werkzeug.serving |
| 22 | +from testing_support.fixture.event_loop import ( # noqa: F401, pylint: disable=W0611 |
| 23 | + event_loop as loop, |
| 24 | +) |
| 25 | +from testing_support.fixtures import ( # noqa: F401, pylint: disable=W0611 |
| 26 | + collector_agent_registration_fixture, |
| 27 | + collector_available_fixture, |
| 28 | +) |
| 29 | + |
| 30 | +PORT = 4443 |
| 31 | +AWS_ACCESS_KEY_ID = "AAAAAAAAAAAACCESSKEY" |
| 32 | +AWS_SECRET_ACCESS_KEY = "AAAAAASECRETKEY" # nosec |
| 33 | +HOST = "127.0.0.1" |
| 34 | + |
| 35 | + |
| 36 | +_default_settings = { |
| 37 | + "transaction_tracer.explain_threshold": 0.0, |
| 38 | + "transaction_tracer.transaction_threshold": 0.0, |
| 39 | + "transaction_tracer.stack_trace_threshold": 0.0, |
| 40 | + "debug.log_data_collector_payloads": True, |
| 41 | + "debug.record_transaction_failure": True, |
| 42 | +} |
| 43 | +collector_agent_registration = collector_agent_registration_fixture( |
| 44 | + app_name="Python Agent Test (external_aiobotocore)", |
| 45 | + default_settings=_default_settings, |
| 46 | + linked_applications=["Python Agent Test (external_aiobotocore)"], |
| 47 | +) |
| 48 | + |
| 49 | + |
| 50 | +def get_free_tcp_port(release_socket: bool = False): |
| 51 | + sckt = socket.socket(socket.AF_INET, socket.SOCK_STREAM) |
| 52 | + sckt.bind((HOST, 0)) |
| 53 | + _, port = sckt.getsockname() # address, port |
| 54 | + if release_socket: |
| 55 | + sckt.close() |
| 56 | + return port |
| 57 | + |
| 58 | + return sckt, port |
| 59 | + |
| 60 | + |
| 61 | +class MotoService: |
| 62 | + """Will Create MotoService. |
| 63 | + Service is ref-counted so there will only be one per process. Real Service will |
| 64 | + be returned by `__aenter__`.""" |
| 65 | + |
| 66 | + _services = {} # {name: instance} |
| 67 | + |
| 68 | + def __init__(self, service_name: str, port: int = None, ssl: bool = False): |
| 69 | + self._service_name = service_name |
| 70 | + |
| 71 | + if port: |
| 72 | + self._socket = None |
| 73 | + self._port = port |
| 74 | + else: |
| 75 | + self._socket, self._port = get_free_tcp_port() |
| 76 | + |
| 77 | + self._thread = None |
| 78 | + self._logger = logging.getLogger("MotoService") |
| 79 | + self._refcount = None |
| 80 | + self._ip_address = HOST |
| 81 | + self._server = None |
| 82 | + self._ssl_ctx = werkzeug.serving.generate_adhoc_ssl_context() if ssl else None |
| 83 | + self._schema = "http" if not self._ssl_ctx else "https" |
| 84 | + |
| 85 | + @property |
| 86 | + def endpoint_url(self): |
| 87 | + return f"{self._schema}://{self._ip_address}:{self._port}" |
| 88 | + |
| 89 | + def __call__(self, func): |
| 90 | + async def wrapper(*args, **kwargs): |
| 91 | + await self._start() |
| 92 | + try: |
| 93 | + result = await func(*args, **kwargs) |
| 94 | + finally: |
| 95 | + await self._stop() |
| 96 | + return result |
| 97 | + |
| 98 | + functools.update_wrapper(wrapper, func) |
| 99 | + wrapper.__wrapped__ = func |
| 100 | + return wrapper |
| 101 | + |
| 102 | + async def __aenter__(self): |
| 103 | + svc = self._services.get(self._service_name) |
| 104 | + if svc is None: |
| 105 | + self._services[self._service_name] = self |
| 106 | + self._refcount = 1 |
| 107 | + await self._start() |
| 108 | + return self |
| 109 | + else: |
| 110 | + svc._refcount += 1 |
| 111 | + return svc |
| 112 | + |
| 113 | + async def __aexit__(self, exc_type, exc_val, exc_tb): |
| 114 | + self._refcount -= 1 |
| 115 | + |
| 116 | + if self._socket: |
| 117 | + self._socket.close() |
| 118 | + self._socket = None |
| 119 | + |
| 120 | + if self._refcount == 0: |
| 121 | + del self._services[self._service_name] |
| 122 | + await self._stop() |
| 123 | + |
| 124 | + def _server_entry(self): |
| 125 | + self._main_app = moto.server.DomainDispatcherApplication( |
| 126 | + moto.server.create_backend_app # , service=self._service_name |
| 127 | + ) |
| 128 | + self._main_app.debug = True |
| 129 | + |
| 130 | + if self._socket: |
| 131 | + self._socket.close() # release right before we use it |
| 132 | + self._socket = None |
| 133 | + |
| 134 | + self._server = werkzeug.serving.make_server( |
| 135 | + self._ip_address, |
| 136 | + self._port, |
| 137 | + self._main_app, |
| 138 | + True, |
| 139 | + ssl_context=self._ssl_ctx, |
| 140 | + ) |
| 141 | + self._server.serve_forever() |
| 142 | + |
| 143 | + async def _start(self): |
| 144 | + self._thread = threading.Thread(target=self._server_entry, daemon=True) |
| 145 | + self._thread.start() |
| 146 | + |
| 147 | + async def _stop(self): |
| 148 | + if self._server: |
| 149 | + self._server.shutdown() |
| 150 | + |
| 151 | + self._thread.join() |
![mergify[bot]](https://avatars.githubusercontent.com/in/10562?v=4&size=40){kind=avatar}
0 commit comments