From 944fd030a2c66cb407a4e83107e021de10c8a199 Mon Sep 17 00:00:00 2001
From: mhorky <mhorky@redhat.com>
Date: Wed, 1 Nov 2023 16:56:21 +0100
Subject: [PATCH] Simplify autoregistration test setup

The function responsible for returning 'server' responses was altered to
be less nested, which should improve its readability.

Cherry-picked from 4d8fcd4.
---
 test/test_auto_registration.py | 95 +++++++++++++++++++---------------
 1 file changed, 52 insertions(+), 43 deletions(-)

diff --git a/test/test_auto_registration.py b/test/test_auto_registration.py
index 983ea2ea4b..5a585845a8 100644
--- a/test/test_auto_registration.py
+++ b/test/test_auto_registration.py
@@ -49,50 +49,59 @@
 AWS_TOKEN = "ABCDEFGHIJKLMNOPQRSTVWXYZabcdefghijklmnopqrstvwxyz0123=="
 
 
-def send_only_imds_v2_is_supported(request, *args, **kwargs):
+def send__aws_imdsv2_only(request, *args, **kwargs):
     """
-    Mock result, when we try to get metadata using GET method against
-    AWS metadata provider. This mock is for the case, when only IMDSv2
-    is supported by instance.
+    Mock result for metadata request on AWS where only IMDSv2 is supported.
+    This function should be used to replace function `requests.Session.send()`.
+
     :param request: HTTP request
-    :return: Mock with result
+    :return: Mocked server result.
     """
-    mock_result = Mock()
-
-    if request.method == 'PUT':
-        if request.url == aws.AWSCloudProvider.CLOUD_PROVIDER_TOKEN_URL:
-            if 'X-aws-ec2-metadata-token-ttl-seconds' in request.headers:
-                mock_result.status_code = 200
-                mock_result.text = AWS_TOKEN
-            else:
-                mock_result.status_code = 400
-                mock_result.text = 'Error: TTL for token not specified'
-        else:
-            mock_result.status_code = 400
-            mock_result.text = 'Error: Invalid URL'
-    elif request.method == 'GET':
-        if 'X-aws-ec2-metadata-token' in request.headers.keys():
-            if request.headers['X-aws-ec2-metadata-token'] == AWS_TOKEN:
-                if request.url == aws.AWSCloudProvider.CLOUD_PROVIDER_METADATA_URL:
-                    mock_result.status_code = 200
-                    mock_result.text = AWS_METADATA
-                elif request.url == aws.AWSCloudProvider.CLOUD_PROVIDER_SIGNATURE_URL:
-                    mock_result.status_code = 200
-                    mock_result.text = AWS_SIGNATURE
-                else:
-                    mock_result.status_code = 400
-                    mock_result.text = 'Error: Invalid URL'
-            else:
-                mock_result.status_code = 400
-                mock_result.text = 'Error: Invalid metadata token provided'
-        else:
-            mock_result.status_code = 400
-            mock_result.text = 'Error: IMDSv1 is not supported on this instance'
-    else:
-        mock_result.status_code = 400
-        mock_result.text = 'Error: not supported request method'
-
-    return mock_result
+    result = Mock()
+
+    if request.method == "PUT":
+        if request.url != aws.AWSCloudProvider.CLOUD_PROVIDER_TOKEN_URL:
+            result.status_code = 400
+            result.text = "Error: Invalid URL"
+            return result
+
+        if "X-aws-ec2-metadata-token-ttl-seconds" not in request.headers:
+            result.status_code = 400
+            result.text = "Error: TTL for token not specified"
+            return result
+
+        result.status_code = 200
+        result.text = AWS_TOKEN
+        return result
+
+    if request.method == "GET":
+        if "X-aws-ec2-metadata-token" not in request.headers.keys():
+            result.status_code = 400
+            result.text = "Error: IMDSv1 is not supported on this instance"
+            return result
+
+        if request.headers["X-aws-ec2-metadata-token"] != AWS_TOKEN:
+            result.status_code = 400
+            result.text = "Error: Invalid metadata token provided"
+            return result
+
+        if request.url == aws.AWSCloudProvider.CLOUD_PROVIDER_METADATA_URL:
+            result.status_code = 200
+            result.text = AWS_METADATA
+            return result
+
+        if request.url == aws.AWSCloudProvider.CLOUD_PROVIDER_SIGNATURE_URL:
+            result.status_code = 200
+            result.text = AWS_SIGNATURE
+            return result
+
+        result.status_code = 400
+        result.text = "Error: Invalid URL"
+        return result
+
+    result.status_code = 400
+    result.text = "Error: not supported request method"
+    return result
 
 
 class TestAutomaticRegistration(unittest.TestCase):
@@ -120,7 +129,7 @@ def test_collect_cloud_info_one_cloud_provider_detected(self):
         one detected cloud provider
         """
         mock_session = Mock()
-        mock_session.send = send_only_imds_v2_is_supported()
+        mock_session.send = send__aws_imdsv2_only
         mock_session.prepare_request = Mock(side_effect=lambda request: request)
         mock_session.hooks = {"response": []}
         aws.AWSCloudProvider._instance._session = mock_session
@@ -150,7 +159,7 @@ def test_collect_cloud_info_more_cloud_providers_detected(self):
         providers were detected
         """
         mock_session = Mock()
-        mock_session.send = send_only_imds_v2_is_supported()
+        mock_session.send = send__aws_imdsv2_only
         mock_session.prepare_request = Mock(side_effect=lambda request: request)
         mock_session.hooks = {"response": []}
         aws.AWSCloudProvider._instance._session = mock_session