From fb4722443d9a5f068bc0c41d0ddd1e514027cdd2 Mon Sep 17 00:00:00 2001
From: Andrei Chekun <andrei.chekun@scylladb.com>
Date: Wed, 30 Oct 2024 10:48:17 +0100
Subject: [PATCH] test.py: Modify s3 server mock

Add the possibility to return environment as a dict to use it later it subprocess created by xdist, without starting another s3 mock server for each thread.
---
 test/pylib/s3_proxy.py       | 16 ++++++++++++++--
 test/pylib/s3_server_mock.py | 16 ++++++++++++++--
 2 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/test/pylib/s3_proxy.py b/test/pylib/s3_proxy.py
index 1be878c08a..622f817001 100644
--- a/test/pylib/s3_proxy.py
+++ b/test/pylib/s3_proxy.py
@@ -220,12 +220,23 @@ class S3ProxyServer:
         self.server.socket.settimeout(10000)
         self.server.socket.listen(1000)
         self.is_running = False
-        os.environ['PROXY_S3_SERVER_PORT'] = f'{port}'
-        os.environ['PROXY_S3_SERVER_HOST'] = host
+        self.envs = {'PROXY_S3_SERVER_PORT': f'{port}', 'PROXY_S3_SERVER_HOST': f'{host}'}
+
+    def _set_environ(self):
+        for key, value in self.envs.items():
+            os.environ[key] = value
+
+    def _unset_environ(self):
+        for key in self.envs.keys():
+            del os.environ[key]
+
+    def get_envs_settings(self):
+        return self.envs
 
     async def start(self):
         if not self.is_running:
             self.logger.info('Starting S3 proxy server on %s', self.server.server_address)
+            self._set_environ()
             loop = asyncio.get_running_loop()
             self.server_thread = loop.run_in_executor(None, self.server.serve_forever)
             self.is_running = True
@@ -233,6 +244,7 @@ class S3ProxyServer:
     async def stop(self):
         if self.is_running:
             self.logger.info('Stopping S3 proxy server')
+            self._unset_environ()
             self.server.shutdown()
             await self.server_thread
             self.is_running = False
diff --git a/test/pylib/s3_server_mock.py b/test/pylib/s3_server_mock.py
index 6e98f500b4..ed9e60bd78 100644
--- a/test/pylib/s3_server_mock.py
+++ b/test/pylib/s3_server_mock.py
@@ -268,12 +268,23 @@ class MockS3Server:
         self.server.socket.listen(1000)
         self.logger = logger
         self.is_running = False
-        os.environ['MOCK_S3_SERVER_PORT'] = f'{port}'
-        os.environ['MOCK_S3_SERVER_HOST'] = host
+        self.envs = {'MOCK_S3_SERVER_PORT': f'{port}', 'MOCK_S3_SERVER_HOST': f'{host}'}
+
+    def _set_environ(self):
+        for key, value in self.envs.items():
+            os.environ[key] = value
+
+    def _unset_environ(self):
+        for key in self.envs.keys():
+            del os.environ[key]
+
+    def get_envs_settings(self):
+        return self.envs
 
     async def start(self):
         if not self.is_running:
             self.logger.info('Starting S3 mock server on %s', self.server.server_address)
+            self._set_environ()
             loop = asyncio.get_running_loop()
             self.server_thread = loop.run_in_executor(None, self.server.serve_forever)
             self.is_running = True
@@ -281,6 +292,7 @@ class MockS3Server:
     async def stop(self):
         if self.is_running:
             self.logger.info('Stopping S3 mock server')
+            self._unset_environ()
             self.server.shutdown()
             await self.server_thread
             self.is_running = False