From 2edf9dd2126dc2dbc230f988abe8443b45885016 Mon Sep 17 00:00:00 2001
From: Jonas Neubert <jonasneu@amazon.com>
Date: Fri, 26 May 2023 08:07:48 -0600
Subject: [PATCH] updates to existing tests

---
 tests/unit/test_args.py   | 12 +++++++++++-
 tests/unit/test_client.py | 15 +++++++++++++--
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/tests/unit/test_args.py b/tests/unit/test_args.py
index 832a689e12..a6a72f36d6 100644
--- a/tests/unit/test_args.py
+++ b/tests/unit/test_args.py
@@ -19,6 +19,7 @@
 from botocore.configprovider import ConfigValueStore
 from botocore.hooks import HierarchicalEmitter
 from botocore.model import ServiceModel
+from botocore.useragent import UserAgentString
 from tests import mock, unittest
 
 
@@ -26,8 +27,15 @@ class TestCreateClientArgs(unittest.TestCase):
     def setUp(self):
         self.event_emitter = mock.Mock(HierarchicalEmitter)
         self.config_store = ConfigValueStore()
+        user_agent_creator = UserAgentString(None, None, None, None, None)
         self.args_create = args.ClientArgsCreator(
-            self.event_emitter, None, None, None, None, self.config_store
+            event_emitter=self.event_emitter,
+            user_agent=None,
+            response_parser_factory=None,
+            loader=None,
+            exceptions_factory=None,
+            config_store=self.config_store,
+            user_agent_creator=user_agent_creator,
         )
         self.service_name = 'ec2'
         self.region = 'us-west-2'
@@ -518,6 +526,7 @@ class TestEndpointResolverBuiltins(unittest.TestCase):
     def setUp(self):
         event_emitter = mock.Mock(HierarchicalEmitter)
         self.config_store = ConfigValueStore()
+        user_agent_creator = UserAgentString(None, None, None, None, None)
         self.args_create = args.ClientArgsCreator(
             event_emitter=event_emitter,
             user_agent=None,
@@ -525,6 +534,7 @@ def setUp(self):
             loader=None,
             exceptions_factory=None,
             config_store=self.config_store,
+            user_agent_creator=user_agent_creator,
         )
         self.bridge = ClientEndpointBridge(
             endpoint_resolver=mock.Mock(),
diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py
index eb7a30787b..c868615b8d 100644
--- a/tests/unit/test_client.py
+++ b/tests/unit/test_client.py
@@ -35,6 +35,7 @@
     UnknownSignatureVersionError,
 )
 from botocore.stub import Stubber
+from botocore.useragent import UserAgentString
 from tests import mock, unittest
 
 
@@ -169,6 +170,7 @@ def create_client_creator(
         endpoint_prefix=None,
         exceptions_factory=None,
         config_store=None,
+        user_agent_creator=None,
     ):
         if event_emitter is None:
             event_emitter = hooks.HierarchicalEmitter()
@@ -187,6 +189,14 @@ def create_client_creator(
             exceptions_factory = ClientExceptionsFactory()
         if config_store is None:
             config_store = self.config_store
+        if user_agent_creator is None:
+            user_agent_creator = (
+                UserAgentString.from_environment().set_session_config(
+                    session_user_agent_name='MyUserAgent',
+                    session_user_agent_version='1.2.3-rc5',
+                    session_user_agent_extra=None,
+                )
+            )
         creator = client.ClientCreator(
             self.loader,
             self.resolver,
@@ -197,6 +207,7 @@ def create_client_creator(
             response_parser_factory,
             exceptions_factory,
             config_store,
+            user_agent_creator,
         )
         return creator
 
@@ -647,7 +658,7 @@ def test_client_user_agent_in_request(self):
             k.lower(): v
             for k, v in self.endpoint.make_request.call_args[0][1].items()
         }
-        self.assertEqual(params['headers']['User-Agent'], 'user-agent')
+        self.assertIn('MyUserAgent/1.2.3', params['headers']['User-Agent'])
 
     def test_client_custom_user_agent_in_request(self):
         creator = self.create_client_creator()
@@ -673,7 +684,7 @@ def test_client_custom_user_agent_extra_in_request(self):
         )
         service_client.test_operation(Foo='one')
         headers = self.endpoint.make_request.call_args[0][1]['headers']
-        self.assertEqual(headers['User-Agent'], 'user-agent extrastuff')
+        self.assertTrue(headers['User-Agent'].endswith('extrastuff'))
 
     def test_client_registers_request_created_handler(self):
         event_emitter = self.create_mock_emitter()