sparkSession in Python

Eu tenho uma função no arquivo principal.py.

def transform_table(table_name, start, end):
  return sparkSession.table(table_name).filter(column.isBetween(start,end))

Estou curioso para saber se a chamada de função sparkSession.table pode ser simulada. Até agora, isso é tudo que tenho… Mensagem de erro não iterável da coluna

@mock.patch("main_file.pyspark.sql.SparkSession")
@pytest.mark.usefixtures("spark")
  def test(self,mocked,spark):
    injected_df = spark.createDataFrame([(1, 100, "ketchup"), (2, 200, "mayo")], ["id", "qty", "condiment"])
    expected_df = spark.createDataFrame([(1, 100, "ketchup"),["id", "qty", "condiment"]) 

    mocked.table.return_value = injected_df
    
    obj = Class(spark)
    obj.transform_table("useless table name", 100, 150)
    assert obj == expected_df

desde já, obrigado

O problema é que está faltando um ]

#errado
expected_df = spark.createDataFrame([(1, 100, "ketchup"),["id", "qty", "condiment"]) 

Mas…

Eu acho que você pode ir mais a fundo e isolar mais o seu teste, não precisar de um ambiente spark.

Você pode substituir o injected_df por um pandas.DataFrame, se você quer algo que se comporte parecido, mas pode não ser 100% compatível. Outra solução é tentar imitar mais a fundo:

mock_df = Mock()
mock_df.somemethod.return_value = "somevalue"
...
mocked.table.return_value = mock_df
...

Ou até um mix dos 2:

df = pandas.DataFrame(...)
mock_df = Mock()
mock_df.somemethod.return_value = df.somemethod(...)
...
mocked.table.return_value = mock_df
...

No fim o objetivo é tentar cumprir o máximo viável de similaridade para testar transform_table