| FazBrowse GitHub Viewer | Trending | | Home |
| Tools: [Download Repo ZIP] [Original HTTPS Page] |
1 parent 6cfdeeb commit 454b441
3 files changed
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -97,10 +97,19 @@ def _get_metadata_ip_root(use_mtls: bool): | |||
| 97 | 97 | # Timeout in seconds to wait for the GCE metadata server when detecting the | |
| 98 | 98 | # GCE environment. | |
| 99 | 99 | try: | |
| 100 | - _METADATA_DEFAULT_TIMEOUT = int(os.getenv("GCE_METADATA_TIMEOUT", 3)) | ||
| 100 | + _METADATA_DEFAULT_TIMEOUT = int(os.getenv(environment_vars.GCE_METADATA_TIMEOUT, 3)) | ||
| 101 | 101 | except ValueError: # pragma: NO COVER | |
| 102 | 102 | _METADATA_DEFAULT_TIMEOUT = 3 | |
| 103 | 103 | ||
| 104 | + # The number of tries to perform when waiting for the GCE metadata server | ||
| 105 | + # when detecting the GCE environment. | ||
| 106 | + try: | ||
| 107 | + _METADATA_DETECT_RETRIES = int( | ||
| 108 | + os.getenv(environment_vars.GCE_METADATA_DETECT_RETRIES, 3) | ||
| 109 | + ) | ||
| 110 | + except ValueError: # pragma: NO COVER | ||
| 111 | + _METADATA_DETECT_RETRIES = 3 | ||
| 112 | + | ||
| 104 | 113 | # This is used to disable checking for the GCE metadata server and directly | |
| 105 | 114 | # assuming it's not available. | |
| 106 | 115 | _NO_GCE_CHECK = os.getenv(environment_vars.NO_GCE_CHECK) == "true" | |
@@ -177,7 +186,9 @@ def _prepare_request_for_mds(request, use_mtls=False) -> None: | |||
| 177 | 186 | request.session.mount(f"https://{host}/", adapter) | |
| 178 | 187 | ||
| 179 | 188 | ||
| 180 | - def ping(request, timeout=_METADATA_DEFAULT_TIMEOUT, retry_count=3): | ||
| 189 | + def ping( | ||
| 190 | + request, timeout=_METADATA_DEFAULT_TIMEOUT, retry_count=_METADATA_DETECT_RETRIES | ||
| 191 | + ): | ||
| 181 | 192 | """Checks to see if the metadata server is available. | |
| 182 | 193 | ||
| 183 | 194 | Args: | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -60,6 +60,16 @@ | |||
| 60 | 60 | """Environment variable providing an alternate ip:port to be used for ip-only | |
| 61 | 61 | GCE metadata requests.""" | |
| 62 | 62 | ||
| 63 | + GCE_METADATA_TIMEOUT = "GCE_METADATA_TIMEOUT" | ||
| 64 | + """Environment variable defining the timeout in seconds to wait for the | ||
| 65 | + GCE metadata server when detecting the GCE environment. | ||
| 66 | + """ | ||
| 67 | + | ||
| 68 | + GCE_METADATA_DETECT_RETRIES = "GCE_METADATA_DETECT_RETRIES" | ||
| 69 | + """Environment variable representing the number of retries that should be | ||
| 70 | + attempted on metadata lookup. | ||
| 71 | + """ | ||
| 72 | + | ||
| 63 | 73 | NO_GCE_CHECK = "NO_GCE_CHECK" | |
| 64 | 74 | """Environment variable controlling whether to check if running on GCE or not. | |
| 65 | 75 | ||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -201,6 +201,23 @@ def test_ping_success_custom_root(mock_metrics_header_value): | |||
| 201 | 201 | ) | |
| 202 | 202 | ||
| 203 | 203 | ||
| 204 | + @mock.patch("google.auth.metrics.mds_ping", return_value=MDS_PING_METRICS_HEADER_VALUE) | ||
| 205 | + def test_ping_failure_custom_retry(mock_metrics_header_value): | ||
| 206 | + request = make_request("") | ||
| 207 | + request.side_effect = exceptions.TransportError() | ||
| 208 | + | ||
| 209 | + os.environ[environment_vars.GCE_METADATA_DETECT_RETRIES] = "10" | ||
| 210 | + importlib.reload(_metadata) | ||
| 211 | + | ||
| 212 | + try: | ||
| 213 | + _metadata.ping(request) | ||
| 214 | + finally: | ||
| 215 | + del os.environ[environment_vars.GCE_METADATA_DETECT_RETRIES] | ||
| 216 | + importlib.reload(_metadata) | ||
| 217 | + | ||
| 218 | + assert request.call_count == 10 | ||
| 219 | + | ||
| 220 | + | ||
| 204 | 221 | def test_get_success_json(): | |
| 205 | 222 | key, value = "foo", "bar" | |
| 206 | 223 | ||
| Back | FazBrowse Home | New Git URL |
0 commit comments