This is an automated email from the ASF dual-hosted git repository.
vincbeck pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 8e26865763 Refactor consolidate import from io in providers (#34378)
8e26865763 is described below
commit 8e2686576399417faf9478d7119110287d4c8630
Author: Miroslav Šedivý <[email protected]>
AuthorDate: Fri Oct 6 18:41:32 2023 +0000
Refactor consolidate import from io in providers (#34378)
---
airflow/providers/amazon/aws/hooks/s3.py | 12 ++++--------
airflow/providers/docker/operators/docker.py | 10 +++++-----
.../amazon/aws/operators/test_s3_file_transform.py | 4 ++--
tests/providers/amazon/aws/operators/test_s3_object.py | 18 +++++++++---------
.../providers/amazon/aws/system/utils/test_helpers.py | 4 ++--
tests/providers/amazon/aws/utils/test_eks_get_token.py | 4 ++--
tests/providers/apache/pinot/hooks/test_pinot.py | 8 ++++----
tests/providers/apache/spark/hooks/test_spark_sql.py | 6 +++---
.../providers/apache/spark/hooks/test_spark_submit.py | 14 +++++++-------
tests/providers/cncf/kubernetes/operators/test_pod.py | 2 +-
.../elasticsearch/log/test_es_task_handler.py | 4 ++--
tests/providers/ftp/hooks/test_ftp.py | 6 +++---
tests/providers/google/cloud/hooks/test_gcs.py | 8 ++++----
.../microsoft/azure/hooks/test_azure_fileshare.py | 4 ++--
tests/providers/sftp/hooks/test_sftp.py | 8 ++++----
tests/providers/ssh/hooks/test_ssh.py | 16 ++++++++--------
tests/system/providers/amazon/aws/example_lambda.py | 16 ++++++++--------
17 files changed, 70 insertions(+), 74 deletions(-)
diff --git a/airflow/providers/amazon/aws/hooks/s3.py
b/airflow/providers/amazon/aws/hooks/s3.py
index f36c671c95..2d6a78f5f6 100644
--- a/airflow/providers/amazon/aws/hooks/s3.py
+++ b/airflow/providers/amazon/aws/hooks/s3.py
@@ -21,7 +21,6 @@ from __future__ import annotations
import asyncio
import fnmatch
import gzip as gz
-import io
import logging
import os
import re
@@ -1120,10 +1119,8 @@ class S3Hook(AwsBaseHook):
if compression == "gzip":
bytes_data = gz.compress(bytes_data)
- file_obj = io.BytesIO(bytes_data)
-
- self._upload_file_obj(file_obj, key, bucket_name, replace, encrypt,
acl_policy)
- file_obj.close()
+ with BytesIO(bytes_data) as f:
+ self._upload_file_obj(f, key, bucket_name, replace, encrypt,
acl_policy)
@unify_bucket_name_and_key
@provide_bucket_name
@@ -1155,9 +1152,8 @@ class S3Hook(AwsBaseHook):
:param acl_policy: The string to specify the canned ACL policy for the
object to be uploaded
"""
- file_obj = io.BytesIO(bytes_data)
- self._upload_file_obj(file_obj, key, bucket_name, replace, encrypt,
acl_policy)
- file_obj.close()
+ with BytesIO(bytes_data) as f:
+ self._upload_file_obj(f, key, bucket_name, replace, encrypt,
acl_policy)
@unify_bucket_name_and_key
@provide_bucket_name
diff --git a/airflow/providers/docker/operators/docker.py
b/airflow/providers/docker/operators/docker.py
index acb8ccb957..94f3be9653 100644
--- a/airflow/providers/docker/operators/docker.py
+++ b/airflow/providers/docker/operators/docker.py
@@ -453,11 +453,11 @@ class DockerOperator(BaseOperator):
# 0 byte file, it can't be anything else than None
return None
# no need to port to a file since we intend to deserialize
- file_standin = BytesIO(b"".join(archived_result))
- tar = tarfile.open(fileobj=file_standin)
- file = tar.extractfile(stat["name"])
- lib = getattr(self, "pickling_library", pickle)
- return lib.loads(file.read())
+ with BytesIO(b"".join(archived_result)) as f:
+ tar = tarfile.open(fileobj=f)
+ file = tar.extractfile(stat["name"])
+ lib = getattr(self, "pickling_library", pickle)
+ return lib.load(file)
try:
return copy_from_docker(self.container["Id"],
self.retrieve_output_path)
diff --git a/tests/providers/amazon/aws/operators/test_s3_file_transform.py
b/tests/providers/amazon/aws/operators/test_s3_file_transform.py
index d147214ed3..49c727b6db 100644
--- a/tests/providers/amazon/aws/operators/test_s3_file_transform.py
+++ b/tests/providers/amazon/aws/operators/test_s3_file_transform.py
@@ -18,10 +18,10 @@
from __future__ import annotations
import errno
-import io
import os
import shutil
import sys
+from io import BytesIO
from tempfile import mkdtemp
from unittest import mock
@@ -39,7 +39,7 @@ class TestS3FileTransformOperator:
self.bucket = "bucket"
self.input_key = "foo"
self.output_key = "bar"
- self.bio = io.BytesIO(self.content)
+ self.bio = BytesIO(self.content)
self.tmp_dir = mkdtemp(prefix="test_tmpS3FileTransform_")
self.transform_script = os.path.join(self.tmp_dir, "transform.py")
os.mknod(self.transform_script)
diff --git a/tests/providers/amazon/aws/operators/test_s3_object.py
b/tests/providers/amazon/aws/operators/test_s3_object.py
index 72e1bdc72a..9507ba5fd3 100644
--- a/tests/providers/amazon/aws/operators/test_s3_object.py
+++ b/tests/providers/amazon/aws/operators/test_s3_object.py
@@ -17,7 +17,7 @@
# under the License.
from __future__ import annotations
-import io
+from io import BytesIO
from unittest import mock
import boto3
@@ -49,7 +49,7 @@ class TestS3CopyObjectOperator:
conn = boto3.client("s3")
conn.create_bucket(Bucket=self.source_bucket)
conn.create_bucket(Bucket=self.dest_bucket)
- conn.upload_fileobj(Bucket=self.source_bucket, Key=self.source_key,
Fileobj=io.BytesIO(b"input"))
+ conn.upload_fileobj(Bucket=self.source_bucket, Key=self.source_key,
Fileobj=BytesIO(b"input"))
# there should be nothing found before S3CopyObjectOperator is executed
assert "Contents" not in conn.list_objects(Bucket=self.dest_bucket,
Prefix=self.dest_key)
@@ -74,7 +74,7 @@ class TestS3CopyObjectOperator:
conn = boto3.client("s3")
conn.create_bucket(Bucket=self.source_bucket)
conn.create_bucket(Bucket=self.dest_bucket)
- conn.upload_fileobj(Bucket=self.source_bucket, Key=self.source_key,
Fileobj=io.BytesIO(b"input"))
+ conn.upload_fileobj(Bucket=self.source_bucket, Key=self.source_key,
Fileobj=BytesIO(b"input"))
# there should be nothing found before S3CopyObjectOperator is executed
assert "Contents" not in conn.list_objects(Bucket=self.dest_bucket,
Prefix=self.dest_key)
@@ -103,7 +103,7 @@ class TestS3DeleteObjectsOperator:
conn = boto3.client("s3")
conn.create_bucket(Bucket=bucket)
- conn.upload_fileobj(Bucket=bucket, Key=key,
Fileobj=io.BytesIO(b"input"))
+ conn.upload_fileobj(Bucket=bucket, Key=key, Fileobj=BytesIO(b"input"))
# The object should be detected before the DELETE action is taken
objects_in_dest_bucket = conn.list_objects(Bucket=bucket, Prefix=key)
@@ -125,7 +125,7 @@ class TestS3DeleteObjectsOperator:
conn = boto3.client("s3")
conn.create_bucket(Bucket=bucket)
for k in keys:
- conn.upload_fileobj(Bucket=bucket, Key=k,
Fileobj=io.BytesIO(b"input"))
+ conn.upload_fileobj(Bucket=bucket, Key=k,
Fileobj=BytesIO(b"input"))
# The objects should be detected before the DELETE action is taken
objects_in_dest_bucket = conn.list_objects(Bucket=bucket,
Prefix=key_pattern)
@@ -147,7 +147,7 @@ class TestS3DeleteObjectsOperator:
conn = boto3.client("s3")
conn.create_bucket(Bucket=bucket)
for k in keys:
- conn.upload_fileobj(Bucket=bucket, Key=k,
Fileobj=io.BytesIO(b"input"))
+ conn.upload_fileobj(Bucket=bucket, Key=k,
Fileobj=BytesIO(b"input"))
# The objects should be detected before the DELETE action is taken
objects_in_dest_bucket = conn.list_objects(Bucket=bucket,
Prefix=key_pattern)
@@ -167,7 +167,7 @@ class TestS3DeleteObjectsOperator:
conn = boto3.client("s3")
conn.create_bucket(Bucket=bucket)
- conn.upload_fileobj(Bucket=bucket, Key=key_of_test,
Fileobj=io.BytesIO(b"input"))
+ conn.upload_fileobj(Bucket=bucket, Key=key_of_test,
Fileobj=BytesIO(b"input"))
# The object should be detected before the DELETE action is tested
objects_in_dest_bucket = conn.list_objects(Bucket=bucket,
Prefix=key_of_test)
@@ -189,7 +189,7 @@ class TestS3DeleteObjectsOperator:
conn = boto3.client("s3")
conn.create_bucket(Bucket=bucket)
- conn.upload_fileobj(Bucket=bucket, Key=key_of_test,
Fileobj=io.BytesIO(b"input"))
+ conn.upload_fileobj(Bucket=bucket, Key=key_of_test,
Fileobj=BytesIO(b"input"))
# The object should be detected before the DELETE action is tested
objects_in_dest_bucket = conn.list_objects(Bucket=bucket,
Prefix=key_of_test)
@@ -235,7 +235,7 @@ class TestS3DeleteObjectsOperator:
conn = boto3.client("s3")
conn.create_bucket(Bucket=bucket)
- conn.upload_fileobj(Bucket=bucket, Key=key_of_test,
Fileobj=io.BytesIO(b"input"))
+ conn.upload_fileobj(Bucket=bucket, Key=key_of_test,
Fileobj=BytesIO(b"input"))
# Set valid values for constructor, and change them later for emulate
rendering template
op = S3DeleteObjectsOperator(
diff --git a/tests/providers/amazon/aws/system/utils/test_helpers.py
b/tests/providers/amazon/aws/system/utils/test_helpers.py
index b977b62362..3db17158da 100644
--- a/tests/providers/amazon/aws/system/utils/test_helpers.py
+++ b/tests/providers/amazon/aws/system/utils/test_helpers.py
@@ -20,9 +20,9 @@ tests/system/providers/amazon/aws/utils/__init__.py
"""
from __future__ import annotations
-import io
import os
import sys
+from io import StringIO
from unittest.mock import ANY, patch
import pytest
@@ -110,7 +110,7 @@ class TestAmazonSystemTestHelpers:
@pytest.mark.parametrize("env_id, is_valid", ENV_ID_TEST_CASES)
def test_validate_env_id_success(self, env_id, is_valid):
if is_valid:
- captured_output = io.StringIO()
+ captured_output = StringIO()
sys.stdout = captured_output
result = _validate_env_id(env_id)
diff --git a/tests/providers/amazon/aws/utils/test_eks_get_token.py
b/tests/providers/amazon/aws/utils/test_eks_get_token.py
index 22de4a0736..c8d66de0db 100644
--- a/tests/providers/amazon/aws/utils/test_eks_get_token.py
+++ b/tests/providers/amazon/aws/utils/test_eks_get_token.py
@@ -17,10 +17,10 @@
from __future__ import annotations
import contextlib
-import io
import json
import os
import runpy
+from io import StringIO
from unittest import mock
from unittest.mock import ANY
@@ -72,7 +72,7 @@ class TestGetEksToken:
mock_eks_hook.return_value.fetch_access_token_for_cluster.return_value
) = "k8s-aws-v1.aHR0cDovL2V4YW1wbGUuY29t"
- with mock.patch("sys.argv", args),
contextlib.redirect_stdout(io.StringIO()) as temp_stdout:
+ with mock.patch("sys.argv", args),
contextlib.redirect_stdout(StringIO()) as temp_stdout:
os.chdir(AIRFLOW_MAIN_FOLDER)
# We are not using run_module because of
https://github.com/pytest-dev/pytest/issues/9007
runpy.run_path("airflow/providers/amazon/aws/utils/eks_get_token.py",
run_name="__main__")
diff --git a/tests/providers/apache/pinot/hooks/test_pinot.py
b/tests/providers/apache/pinot/hooks/test_pinot.py
index 89dde59666..6c60c5d6ec 100644
--- a/tests/providers/apache/pinot/hooks/test_pinot.py
+++ b/tests/providers/apache/pinot/hooks/test_pinot.py
@@ -17,9 +17,9 @@
# under the License.
from __future__ import annotations
-import io
import os
import subprocess
+from io import BytesIO
from unittest import mock
import pytest
@@ -158,7 +158,7 @@ class TestPinotAdminHook:
def test_run_cli_success(self, mock_popen):
mock_proc = mock.MagicMock()
mock_proc.returncode = 0
- mock_proc.stdout = io.BytesIO(b"")
+ mock_proc.stdout = BytesIO(b"")
mock_popen.return_value.__enter__.return_value = mock_proc
params = ["foo", "bar", "baz"]
@@ -173,7 +173,7 @@ class TestPinotAdminHook:
msg = b"Exception caught"
mock_proc = mock.MagicMock()
mock_proc.returncode = 0
- mock_proc.stdout = io.BytesIO(msg)
+ mock_proc.stdout = BytesIO(msg)
mock_popen.return_value.__enter__.return_value = mock_proc
params = ["foo", "bar", "baz"]
with pytest.raises(AirflowException):
@@ -187,7 +187,7 @@ class TestPinotAdminHook:
def test_run_cli_failure_status_code(self, mock_popen):
mock_proc = mock.MagicMock()
mock_proc.returncode = 1
- mock_proc.stdout = io.BytesIO(b"")
+ mock_proc.stdout = BytesIO(b"")
mock_popen.return_value.__enter__.return_value = mock_proc
self.db_hook.pinot_admin_system_exit = True
diff --git a/tests/providers/apache/spark/hooks/test_spark_sql.py
b/tests/providers/apache/spark/hooks/test_spark_sql.py
index 9bd46e6ce3..ee0f4f359c 100644
--- a/tests/providers/apache/spark/hooks/test_spark_sql.py
+++ b/tests/providers/apache/spark/hooks/test_spark_sql.py
@@ -17,8 +17,8 @@
# under the License.
from __future__ import annotations
-import io
import itertools
+from io import StringIO
from unittest.mock import call, patch
import pytest
@@ -85,8 +85,8 @@ class TestSparkSqlHook:
@patch("airflow.providers.apache.spark.hooks.spark_sql.subprocess.Popen")
def test_spark_process_runcmd(self, mock_popen):
# Given
- mock_popen.return_value.stdout = io.StringIO("Spark-sql communicates
using stdout")
- mock_popen.return_value.stderr = io.StringIO("stderr")
+ mock_popen.return_value.stdout = StringIO("Spark-sql communicates
using stdout")
+ mock_popen.return_value.stderr = StringIO("stderr")
mock_popen.return_value.wait.return_value = 0
# When
diff --git a/tests/providers/apache/spark/hooks/test_spark_submit.py
b/tests/providers/apache/spark/hooks/test_spark_submit.py
index 719e4383c4..956f0cd6cb 100644
--- a/tests/providers/apache/spark/hooks/test_spark_submit.py
+++ b/tests/providers/apache/spark/hooks/test_spark_submit.py
@@ -17,8 +17,8 @@
# under the License.
from __future__ import annotations
-import io
import os
+from io import StringIO
from unittest.mock import call, patch
import pytest
@@ -241,8 +241,8 @@ class TestSparkSubmitHook:
@patch("airflow.providers.apache.spark.hooks.spark_submit.subprocess.Popen")
def test_spark_process_runcmd(self, mock_popen):
# Given
- mock_popen.return_value.stdout = io.StringIO("stdout")
- mock_popen.return_value.stderr = io.StringIO("stderr")
+ mock_popen.return_value.stdout = StringIO("stdout")
+ mock_popen.return_value.stderr = StringIO("stderr")
mock_popen.return_value.wait.return_value = 0
# When
@@ -694,8 +694,8 @@ class TestSparkSubmitHook:
@patch("airflow.providers.apache.spark.hooks.spark_submit.subprocess.Popen")
def test_yarn_process_on_kill(self, mock_popen, mock_renew_from_kt):
# Given
- mock_popen.return_value.stdout = io.StringIO("stdout")
- mock_popen.return_value.stderr = io.StringIO("stderr")
+ mock_popen.return_value.stdout = StringIO("stdout")
+ mock_popen.return_value.stderr = StringIO("stderr")
mock_popen.return_value.poll.return_value = None
mock_popen.return_value.wait.return_value = 0
log_lines = [
@@ -776,8 +776,8 @@ class TestSparkSubmitHook:
@patch("airflow.providers.apache.spark.hooks.spark_submit.subprocess.Popen")
def test_k8s_process_on_kill(self, mock_popen, mock_client_method):
# Given
- mock_popen.return_value.stdout = io.StringIO("stdout")
- mock_popen.return_value.stderr = io.StringIO("stderr")
+ mock_popen.return_value.stdout = StringIO("stdout")
+ mock_popen.return_value.stderr = StringIO("stderr")
mock_popen.return_value.poll.return_value = None
mock_popen.return_value.wait.return_value = 0
client = mock_client_method.return_value
diff --git a/tests/providers/cncf/kubernetes/operators/test_pod.py
b/tests/providers/cncf/kubernetes/operators/test_pod.py
index 2af333dcff..d804f51469 100644
--- a/tests/providers/cncf/kubernetes/operators/test_pod.py
+++ b/tests/providers/cncf/kubernetes/operators/test_pod.py
@@ -18,6 +18,7 @@ from __future__ import annotations
import re
from contextlib import contextmanager, nullcontext
+from io import BytesIO
from unittest import mock
from unittest.mock import MagicMock, patch
@@ -25,7 +26,6 @@ import pendulum
import pytest
from kubernetes.client import ApiClient, V1PodSecurityContext, V1PodStatus,
models as k8s
from urllib3 import HTTPResponse
-from urllib3.packages.six import BytesIO
from airflow.exceptions import AirflowException, AirflowSkipException,
TaskDeferred
from airflow.models import DAG, DagModel, DagRun, TaskInstance
diff --git a/tests/providers/elasticsearch/log/test_es_task_handler.py
b/tests/providers/elasticsearch/log/test_es_task_handler.py
index f720429825..def1c0ab45 100644
--- a/tests/providers/elasticsearch/log/test_es_task_handler.py
+++ b/tests/providers/elasticsearch/log/test_es_task_handler.py
@@ -17,12 +17,12 @@
# under the License.
from __future__ import annotations
-import io
import json
import logging
import os
import re
import shutil
+from io import StringIO
from pathlib import Path
from unittest import mock
from urllib.parse import quote
@@ -602,7 +602,7 @@ class TestElasticsearchTaskHandler:
self.es_task_handler.frontend = frontend
assert self.es_task_handler.supports_external_link == expected
- @mock.patch("sys.__stdout__", new_callable=io.StringIO)
+ @mock.patch("sys.__stdout__", new_callable=StringIO)
def test_dynamic_offset(self, stdout_mock, ti, time_machine):
# arrange
handler = ElasticsearchTaskHandler(
diff --git a/tests/providers/ftp/hooks/test_ftp.py
b/tests/providers/ftp/hooks/test_ftp.py
index dcb2e4ac51..fd810aa604 100644
--- a/tests/providers/ftp/hooks/test_ftp.py
+++ b/tests/providers/ftp/hooks/test_ftp.py
@@ -17,7 +17,7 @@
# under the License.
from __future__ import annotations
-import io
+from io import StringIO
from unittest import mock
from airflow.providers.ftp.hooks import ftp as fh
@@ -107,14 +107,14 @@ class TestFTPHook:
self.conn_mock.size.assert_called_once_with(path)
def test_retrieve_file(self):
- _buffer = io.StringIO("buffer")
+ _buffer = StringIO("buffer")
with fh.FTPHook() as ftp_hook:
ftp_hook.retrieve_file(self.path, _buffer)
self.conn_mock.retrbinary.assert_called_once_with("RETR path",
_buffer.write, 8192)
def test_retrieve_file_with_callback(self):
func = mock.Mock()
- _buffer = io.StringIO("buffer")
+ _buffer = StringIO("buffer")
with fh.FTPHook() as ftp_hook:
ftp_hook.retrieve_file(self.path, _buffer, callback=func)
self.conn_mock.retrbinary.assert_called_once_with("RETR path", func,
8192)
diff --git a/tests/providers/google/cloud/hooks/test_gcs.py
b/tests/providers/google/cloud/hooks/test_gcs.py
index f3c17e9d42..5a22e0a464 100644
--- a/tests/providers/google/cloud/hooks/test_gcs.py
+++ b/tests/providers/google/cloud/hooks/test_gcs.py
@@ -18,11 +18,11 @@
from __future__ import annotations
import copy
-import io
import logging
import os
import re
from datetime import datetime, timedelta
+from io import BytesIO
from unittest import mock
import dateutil
@@ -699,7 +699,7 @@ class TestGCSHook:
def test_download_as_bytes(self, mock_service):
test_bucket = "test_bucket"
test_object = "test_object"
- test_object_bytes = io.BytesIO(b"input")
+ test_object_bytes = BytesIO(b"input")
download_method =
mock_service.return_value.bucket.return_value.blob.return_value.download_as_bytes
download_method.return_value = test_object_bytes
@@ -713,7 +713,7 @@ class TestGCSHook:
def test_download_to_file(self, mock_service):
test_bucket = "test_bucket"
test_object = "test_object"
- test_object_bytes = io.BytesIO(b"input")
+ test_object_bytes = BytesIO(b"input")
test_file = "test_file"
download_filename_method = (
@@ -737,7 +737,7 @@ class TestGCSHook:
def test_provide_file(self, mock_service, mock_temp_file):
test_bucket = "test_bucket"
test_object = "test_object"
- test_object_bytes = io.BytesIO(b"input")
+ test_object_bytes = BytesIO(b"input")
test_file = "test_file"
download_filename_method = (
diff --git a/tests/providers/microsoft/azure/hooks/test_azure_fileshare.py
b/tests/providers/microsoft/azure/hooks/test_azure_fileshare.py
index 8456cd3602..819209ccb3 100644
--- a/tests/providers/microsoft/azure/hooks/test_azure_fileshare.py
+++ b/tests/providers/microsoft/azure/hooks/test_azure_fileshare.py
@@ -17,7 +17,7 @@
# under the License.
from __future__ import annotations
-import io
+from io import StringIO
from unittest import mock
import pytest
@@ -139,7 +139,7 @@ class TestAzureFileshareHook:
hook = AzureFileShareHook(
azure_fileshare_conn_id="azure_fileshare_extras",
share_name="share", file_path="file"
)
- data = io.StringIO("stream")
+ data = StringIO("stream")
hook.get_file_to_stream(stream=data)
mock_instance.download_file.assert_called_once_with()
diff --git a/tests/providers/sftp/hooks/test_sftp.py
b/tests/providers/sftp/hooks/test_sftp.py
index b3f1f5dbb8..a1f6f930db 100644
--- a/tests/providers/sftp/hooks/test_sftp.py
+++ b/tests/providers/sftp/hooks/test_sftp.py
@@ -34,10 +34,10 @@ from airflow.utils.session import provide_session
def generate_host_key(pkey: paramiko.PKey):
- key_fh = StringIO()
- pkey.write_private_key(key_fh)
- key_fh.seek(0)
- key_obj = paramiko.RSAKey(file_obj=key_fh)
+ with StringIO() as key_fh:
+ pkey.write_private_key(key_fh)
+ key_fh.seek(0)
+ key_obj = paramiko.RSAKey(file_obj=key_fh)
return key_obj.get_base64()
diff --git a/tests/providers/ssh/hooks/test_ssh.py
b/tests/providers/ssh/hooks/test_ssh.py
index 352b6692da..14717b904d 100644
--- a/tests/providers/ssh/hooks/test_ssh.py
+++ b/tests/providers/ssh/hooks/test_ssh.py
@@ -48,18 +48,18 @@ conn.sendall(b'hello')
def generate_key_string(pkey: paramiko.PKey, passphrase: str | None = None):
- key_fh = StringIO()
- pkey.write_private_key(key_fh, password=passphrase)
- key_fh.seek(0)
- key_str = key_fh.read()
+ with StringIO() as key_fh:
+ pkey.write_private_key(key_fh, password=passphrase)
+ key_fh.seek(0)
+ key_str = key_fh.read()
return key_str
def generate_host_key(pkey: paramiko.PKey):
- key_fh = StringIO()
- pkey.write_private_key(key_fh)
- key_fh.seek(0)
- key_obj = paramiko.RSAKey(file_obj=key_fh)
+ with StringIO() as key_fh:
+ pkey.write_private_key(key_fh)
+ key_fh.seek(0)
+ key_obj = paramiko.RSAKey(file_obj=key_fh)
return key_obj.get_base64()
diff --git a/tests/system/providers/amazon/aws/example_lambda.py
b/tests/system/providers/amazon/aws/example_lambda.py
index 39ca80b5c9..4ba74a26d5 100644
--- a/tests/system/providers/amazon/aws/example_lambda.py
+++ b/tests/system/providers/amazon/aws/example_lambda.py
@@ -16,10 +16,10 @@
# under the License.
from __future__ import annotations
-import io
import json
import zipfile
from datetime import datetime
+from io import BytesIO
import boto3
@@ -49,13 +49,13 @@ def test(*args):
# Create a zip file containing one file "lambda_function.py" to deploy to the
lambda function
def create_zip(content: str):
- zip_output = io.BytesIO()
- with zipfile.ZipFile(zip_output, "w", zipfile.ZIP_DEFLATED) as zip_file:
- info = zipfile.ZipInfo("lambda_function.py")
- info.external_attr = 0o777 << 16
- zip_file.writestr(info, content)
- zip_output.seek(0)
- return zip_output.read()
+ with BytesIO() as zip_output:
+ with zipfile.ZipFile(zip_output, "w", zipfile.ZIP_DEFLATED) as
zip_file:
+ info = zipfile.ZipInfo("lambda_function.py")
+ info.external_attr = 0o777 << 16
+ zip_file.writestr(info, content)
+ zip_output.seek(0)
+ return zip_output.read()
@task(trigger_rule=TriggerRule.ALL_DONE)
