Невозможно заменить определение функции Python на макет и unitest ⇐ Python

[Anonymous]

Я пытаюсь написать интеграционный тест для следующего кода Python:

Код: Выделить всё

import xx.settings.config as stg
from xx.infrastructure.utils import csvReader, dataframeWriter
from pyspark.sql import SparkSession
from typing import List
from awsglue.utils import getResolvedOptions
import sys

def main(argv: List[str]) -> None:
args = getResolvedOptions(
argv,
['JOB_NAME', 'S3_BRONZE_BUCKET_NAME', 'S3_PRE_SILVER_BUCKET_NAME', 'S3_BRONZE_PATH', 'S3_PRE_SILVER_PATH'],
)

s3_bronze_bucket_name = args['S3_BRONZE_BUCKET_NAME']
s3_pre_silver_bucket_name = args['S3_PRE_SILVER_BUCKET_NAME']
s3_bronze_path = args['S3_BRONZE_PATH']
s3_pre_silver_path = args['S3_PRE_SILVER_PATH']

spark = SparkSession.builder.getOrCreate()
spark.conf.set('spark.sql.sources.partitionOverwriteMode', 'dynamic')

for table in list(stg.data_schema.keys()):
raw_data = stg.data_schema[table].columns.to_dict()
df = csvReader(spark, s3_bronze_bucket_name, s3_bronze_path, table, schema, '\t')
dataframeWriter(df, s3_pre_silver_bucket_name, s3_pre_silver_path, table, stg.data_schema[table].partitionKey)

if __name__ == '__main__':
main(sys.argv)

По сути, я просматриваю список таблиц, затем считываю их содержимое (формат csv) из S3 и также записываю их в формате паркета в S3.
Это определения csvReader и dataframeWriter:

Код: Выделить всё

def csvReader(spark: SparkSession, bucket: str, path: str, table: str, schema: StructType, sep: str) -> DataFrame:
return (
spark.read.format('csv')
.option('header', 'true')
.option('sep', sep)
.schema(schema)
.load(f's3a://{bucket}/{path}/{table}.csv')
)

def dataframeWriter(df: DataFrame, bucket: str, path: str, table: str, partition_key: str) -> None:
df.write.partitionBy(partition_key).mode('overwrite').parquet(f's3a://{bucket}/{path}/{table}/')

Для моих интеграционных тестов я хотел бы заменить взаимодействие S3 на
взаимодействие с локальными файлами (прочитайте css из локального и напишите паркет локально. Вот что я сделал:

Код: Выделить всё

import os
from unittest import TestCase
from unittest.mock import patch, Mock
from pyspark.sql import SparkSession
from pyspark.sql.types import StructType
import xx.application.perfmarket_pre_silver as perfmarket_pre_silver
from dvr_config_utils.config import initialize_settings

def local_csvReader(spark: SparkSession, table: str, schema: StructType, sep: str):
"""Mocked function that replaces real csvReader. this one reads from local rather than S3."""
return (
spark.read.format('csv')
.option('header', 'true')
.option('sep', sep)
.schema(schema)
.load(f'../input_mock/{table}.csv')
)

def local_dataframeWriter(df, table: str, partition_key: str):
"""Mocked function that replaces real dataframeWriter.  this one writes in local rather than S3."""
output_dir = f'../output_mock/{table}/'
if not os.path.exists(output_dir):
os.makedirs(output_dir)
df.write.partitionBy(partition_key).mode('overwrite').parquet(output_dir)

class TestPerfmarketSilver(TestCase):
@classmethod
def setUpClass(cls):
cls.spark = SparkSession.builder.master('local').appName('TestPerfmarketSilver').getOrCreate()
cls.spark.conf.set('spark.sql.sources.partitionOverwriteMode', 'dynamic')

@classmethod
def tearDownClass(cls):
"""Clean up the Spark session and test data."""
cls.spark.stop()
os.system('rm -rf ../output_mock')

@patch('xx.application.cc.getResolvedOptions')
@patch('src.xx.infrastructure.utils.csvReader', side_effect=local_csvReader)
@patch('xx.infrastructure.utils.dataframeWriter', side_effect=local_dataframeWriter)
def test_main(self, mock_csvreader, mock_datawriter, mocked_get_resolved_options: Mock):

expected_results = {'chemins': {'nbRows': 8}}

mocked_get_resolved_options.return_value = {
'JOB_NAME': 'perfmarket_pre_silver_test',
'S3_BRONZE_BUCKET_NAME': 'test_bronze',
'S3_PRE_SILVER_BUCKET_NAME': 'test_pre_silver',
'S3_BRONZE_PATH': '../input_mock',
'S3_PRE_SILVER_PATH': '../output_mock'
}
perfmarket_pre_silver.main([])

for table in stg.data_schema.keys():
# Verify that the output Parquet file is created
output_path = f'../output_mock/{table}/'
self.assertTrue(os.path.exists(output_path))

# Read the written Parquet file and check the data
written_df = self.spark.read.parquet(output_path)
self.assertEqual(written_df.count(), expected_results[table]['nbRows'])  # Check row count
self.assertTrue(
[
column_data['bronze_name']
for table in stg.data_schema.values()
for column_data in table['columns'].values()
]
== written_df.columns
)

Что я хотел сделать с этими двумя строками:

Код: Выделить всё

@patch('src.xx.infrastructure.utils.csvReader', side_effect=local_csvReader)
@patch('xx.infrastructure.utils.dataframeWriter', side_effect=local_dataframeWriter)

Заменить определения csvReader на local_csvReader и dataframeWriter на local_dataframeWriter.
К сожалению, код перенастраивается

Код: Выделить всё

py4j.protocol.Py4JJavaError: An error occurred while calling o39.load.
: java.lang.RuntimeException: java.lang.ClassNotFoundException: Class org.apache.hadoop.fs.s3a.S3AFileSystem not found

Ошибка указывает на вызов csvReader в основном коде (первый фрагмент).
Так что моя техника замены явно не работает.Что я делаю не так?

Подробнее здесь: https://stackoverflow.com/questions/790 ... nd-unitest