単体テスト

Copied!
1
2
3
4
5
6
7
8
import duckdb

def test_dataframe():
    conn = duckdb.connect()
    df = conn.sql("""SELECT * FROM VALUES
                    ('a', 1), ('b', 2)
                    AS t(letter, number)""").fetchdf()
    assert list(df.columns) == ['letter', 'number']

Copied!
1
2
3
4
5
import pandas as pd

def test_dataframe():
    df = pd.DataFrame([['a', 1], ['b', 2]], columns=['letter', 'number'])
    assert list(df.columns) == ['letter', 'number']

Copied!
1
2
3
def test_dataframe(spark_session):
    df = spark_session.createDataFrame([['a', 1], ['b', 2]], ['letter', 'number'])
    assert df.schema.names == ['letter', 'number']

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
from transforms.api import transform, Input, Output


@transform.using(
    output=Output("<output_dataset_rid>"),
    aircraft=Input("<input_dataset_rid>"),
)
def compute(ctx, output, aircraft):
    conn = ctx.duckdb().conn
    filtered_query = conn.sql("""
        SELECT * FROM aircraft
        WHERE number_of_seats > 300
        AND operating_status = 'Yes'
    """)
    output.write_table(filtered_query)

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
from transforms.api import transform, Input, Output
import pandas as pd


@transform.using(
    output=Output("<output_dataset_rid>"),
    aircraft=Input("<input_dataset_rid>"),
)
def compute(output, aircraft):
    aircraft_df = aircraft.pandas()
    filtered_df = aircraft_df[
        (aircraft_df["number_of_seats"] > 300) &
        (aircraft_df["operating_status"] == "Yes")
    ]
    output.write_table(filtered_df)

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
from pyspark.sql import functions as F
from transforms.api import transform, Input, Output


@transform.spark.using(
    output=Output("<output_dataset_rid>"),
    aircraft=Input("<input_dataset_rid>"),
)
def compute(output, aircraft):
    aircraft_df = aircraft.dataframe()
    filtered_df = aircraft_df.filter((F.col("number_of_seats") > F.lit(300)) & (F.col("operating_status") == F.lit("Yes")))

    output.write_dataframe(filtered_df)

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
import os
from pathlib import Path
import duckdb
import pandas as pd
from unittest.mock import MagicMock, patch
from myproject.datasets.find_aircraft import compute

TEST_RESOURCES_DIRECTORY_PATH = Path(os.path.dirname(__file__)).joinpath('resources')


def test_find_aircrafts():
    conn = duckdb.connect()

    # DuckDB を使用して CSV データを読み込みます
    aircraft_mock_df = conn.sql(f"SELECT * FROM read_csv('{TEST_RESOURCES_DIRECTORY_PATH / 'aircraft_mock.csv'}', header=true)").fetchdf()
    expected_filtered_aircraft_df = conn.sql(f"SELECT * FROM read_csv('{TEST_RESOURCES_DIRECTORY_PATH / 'expected_filtered_aircraft.csv'}', header=true)").fetchdf()

    # 入力と出力のモックオブジェクトを作成します
    results_output_mock = MagicMock()
    aircraft_mock_input = MagicMock()

    # DuckDB 接続を持つモックコンテキストを作成します
    ctx_mock = MagicMock()
    duckdb_mock = MagicMock()
    duckdb_mock.conn = conn
    ctx_mock.duckdb = MagicMock(return_value=duckdb_mock)

    # DuckDB 接続にテストデータを登録します
    conn.register('aircraft', aircraft_mock_df)

    # データセットの入出力をモック化します
    datasets = {"results_output": results_output_mock, "aircraft_input": aircraft_mock_input}
    module = MagicMock()
    with patch.dict("sys.modules", {"foundry": MagicMock(), "foundry.transforms": module}):
        module.Dataset.get.side_effect = datasets.get
        compute(ctx_mock, results_output_mock, aircraft_mock_input)

    # モック化した write_table 呼び出しから結果を取得します
    result_query = results_output_mock.write_table.call_args[0][0]
    result_df = result_query.fetchdf()

    assert list(result_df.columns) == list(expected_filtered_aircraft_df.columns)
    assert len(result_df) == len(expected_filtered_aircraft_df)
    pd.testing.assert_frame_equal(
        result_df.sort_values(by=list(result_df.columns)).reset_index(drop=True),
        expected_filtered_aircraft_df.sort_values(by=list(expected_filtered_aircraft_df.columns)).reset_index(drop=True),
        check_dtype=False
    )

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
import os
from pathlib import Path
import pandas as pd
from myproject.datasets.find_aircraft import compute

TEST_RESOURCES_DIRECTORY_PATH = Path(os.path.dirname(__file__)).joinpath('resources')


def test_find_aircrafts():
    aircraft_mock_df = pd.read_csv(TEST_RESOURCES_DIRECTORY_PATH / 'aircraft_mock.csv')
    expected_filtered_aircraft_df = pd.read_csv(TEST_RESOURCES_DIRECTORY_PATH / 'expected_filtered_aircraft.csv')

    # 入力と出力のモックオブジェクトを作成します
    results_output_mock = MagicMock()
    aircraft_mock_input = MagicMock()
    aircraft_mock_input.pandas = MagicMock(return_value=aircraft_mock_df)
    aircraft_mock_input.read_table = MagicMock(return_value=aircraft_mock_df)

    # データセットの入出力をモック化します
    datasets = {"results_output": results_output_mock, "aircraft_input": aircraft_mock_input}
    module = MagicMock()
    with patch.dict("sys.modules", {"foundry": MagicMock(), "foundry.transforms": module}):
        module.Dataset.get.side_effect = datasets.get
        compute()

    result_df = results_output_mock.write_pandas.call_args[0][0]

    assert list(result_df.columns) == list(expected_filtered_aircraft_df.columns)
    assert len(result_df) == len(expected_filtered_aircraft_df)
    pd.testing.assert_frame_equal(
        result_df.sort_values(by=list(result_df.columns)).reset_index(drop=True),
        expected_filtered_aircraft_df.sort_values(by=list(expected_filtered_aircraft_df.columns)).reset_index(drop=True),
        check_dtype=False
    )

Copied!
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
import os
from pathlib import Path
from unittest.mock import MagicMock
from transforms.api import Input
from myproject.datasets.find_aircraft import compute

TEST_RESOURCES_DIRECTORY_PATH = Path(os.path.dirname(__file__)).joinpath('resources')


def test_find_aircrafts(spark_session):
    aircraft_mock_df = spark_session.read.csv(
       str(TEST_RESOURCES_DIRECTORY_PATH.joinpath('aircraft_mock.csv')),
       inferSchema=True,
       header=True
        )

    expected_filtered_aircraft_df = spark_session.read.csv(
       str(TEST_RESOURCES_DIRECTORY_PATH.joinpath('expected_filtered_aircraft.csv')),
       inferSchema=True,
       header=True
        )

    # 出力のモックオブジェクトを作成します
    results_output_mock = MagicMock()

    # 入力のラッパーを作成し、返されるデータセットを設定します
    aircraft_mock_input = Input()
    aircraft_mock_input.dataframe = MagicMock(return_value=aircraft_mock_df)

    # モック出力オブジェクトを使用してトランスフォームを実行します
    compute(
        results_output=results_output_mock,
        aircraft_input=aircraft_mock_input
    )

    # write_dataframe の呼び出しに使用された引数を取得し、
    # トランスフォームの結果として書き込み用に送信された DataFrame を抽出します
    args, kwargs = results_output_mock.write_dataframe.call_args
    result_df = args[0]

    assert result_df.columns == expected_filtered_aircraft_df.columns
    assert result_df.count() == expected_filtered_aircraft_df.count()
    assert result_df.exceptAll(expected_filtered_aircraft_df).count() == 0
    assert expected_filtered_aircraft_df.exceptAll(result_df).count() == 0