Testes unitários de pipelines

Importante

Este recurso está em versão Beta.

Para informações gerais sobre testes unitários em Python em Databricks, veja Python unit testing.

Os pipelines Lakeflow suportam a escrita de testes unitários em Python no editor web Lakeflow Pipelines. Isto permite-lhe validar a lógica de transformação em Python ou SQL usando dados simulados. Com o framework de testes de pipeline, pode testar casos limites, validar APIs proprietárias de pipeline (Auto CDC, tabelas de streaming, expectativas, fluxos anexados) e iterar usando entradas simuladas para operações suportadas com identificadores de tabela. Revise as limitações de isolamento antes de realizar testes.

  • Execução isolada de testes: O framework fornece um SparkSession que redireciona as operações das tabelas para um esquema de teste temporário no catálogo predefinido do pipeline, para que possa simular dados de entrada e escrever resultados de teste sem afetar as tabelas de produção. O isolamento aplica-se a operações que referenciam uma tabela pelo nome; ver Limitações.
  • Âmbito de teste flexível: Executar um subconjunto de um pipeline (tabelas individuais, cadeias de tabelas dependentes ou pipelines inteiros) no cálculo do pipeline usando o teste SparkSession.
  • Validação de resultados: Verificar os resultados das tabelas de saída isoladas criadas num teste usando asserções pytest padrão.

Quando usar testes unitários

Os casos de uso típicos incluem:

  • Validação de nova lógica de transformação: Teste se a sua transformação produz o esquema esperado, contagens de linhas, agregações e lógica de negócio antes de ser executada com dados de produção.
  • Testar as especificações do Auto CDC: Valide que as definições de fluxo do Auto CDC processam corretamente os eventos de alteração, lidando com inserções, atualizações, eliminações e tipos SCD (Slowly Changing Dimension), usando dados simulados.
  • Testes de expectativas e regras de qualidade dos dados: Verifique se as expectativas falham quando devem e passam quando os dados são válidos.
  • Testar entre tabelas dependentes: Teste cadeias de transformações (por exemplo, bronze, prata e ouro) para validar se os dados fluem corretamente ao longo do grafo do pipeline.

Requirements

  • Permissão de Ownerpipeline, mais os privilégios USE CATALOG e CREATE SCHEMA no catálogo padrão do pipeline. O framework necessita destes privilégios para criar o esquema de teste temporário onde os testes são executados.

    Para verificar ou definir a permissão do pipeline, abra o pipeline e clique em Partilhar. Deve ser o pipeline Owner (IS OWNER); CAN RUN e CAN MANAGE não são suficientes para executar testes. Veja Configurar permissões de pipeline.

    Para verificar ou definir os privilégios de catálogo, abra o catálogo no Explorador de Catálogos, selecione o separador Permissões e confirme que tem USE CATALOG e CREATE SCHEMA. Um proprietário de catálogo, um administrador de metastore ou um utilizador com esse MANAGE privilégio podem concedê-los, incluindo com SQL:

    GRANT USE CATALOG, CREATE SCHEMA ON CATALOG <catalog_name> TO `<principal>`;

    Para mais informações, consulte a referência de privilégios do Unity Catalog.

  • O pipeline deve ser configurado em modo desencadeado (não contínuo).

  • O pipeline tem de estar no canal PREVIEW. Os testes unitários estão em versão beta e só estão disponíveis no PREVIEW.

  • O Spark Connect não é suportado.

Note

O isolamento de testes abrange operações de tabela que referenciam uma tabela pelo nome. Operações que contornam o isolamento podem ocorrer tanto no seu código de teste como em qualquer código de pipeline executado pelas saídas que selecionar, incluindo as suas dependências transitivas. Um ficheiro de teste que parece seguro pode, ainda assim, executar um fluxo de um pipeline que lê ou escreve através de um caminho ou de um conector, e que opera sobre dados de produção. Para evitar que os testes afetem dados ou metadados de produção, siga estas regras:

  • Referenciar cada tabela pelo nome (catalog.schema.table), e simular todas as entradas pelo nome. Não leia nem escreva utilizando o caminho (/Volumes/..., dbfs:/..., s3://..., abfss://...) e não leia a partir de conectores como o Kafka ou o Auto Loader. Estas contornam o isolamento e atuam sobre sistemas reais de produção.
  • Não faça declarações de governação ou propriedade, como GRANT, REVOKE, ALTER ... OWNER TO, SET/UNSET TAGS, ou CREATE/DROP POLICY. Estes são executados sobre o securável real de produção.
  • Não crie catálogos ou esquemas (CREATE CATALOG, CREATE SCHEMA). Estes chegam à sua verdadeira metastore do Unity Catalog.
  • Não execute todo o fluxo de processamento se o respetivo grafo incluir entradas baseadas em caminhos de ficheiro, conectores, operações de escrita imperativas ou outros efeitos secundários externos. Selecione apenas as saídas cujas dependências utilizam operações de tabela de catálogo suportadas e foram substituídas por entradas simuladas.

Consulte Limitações para obter detalhes.

Limitações

Warning

Algumas operações contornam o isolamento de testes e podem atuar sobre dados reais de produção ou metadados. Revise as seguintes limitações antes de realizar os testes.

O isolamento dos testes baseia-se apenas no nome da tabela

  • Não leia nem escreva através do caminho ou do conector. O isolamento redireciona apenas operações que referenciam uma tabela pelo nome (por exemplo, spark.read.table("catalog.schema.table") ou df.write.saveAsTable("catalog.schema.table")). As operações efetuadas por um caminho ou através de um conector contornam o isolamento e atuam diretamente em sistemas de produção reais:

    • Escrever através de um caminho (por exemplo, df.write.save("/Volumes/..."), um caminho dbfs:/, ou um caminho na cloud ou de localização externa, como s3://... ou abfss://...) escreve no armazenamento de produção real e pode sobrescrever dados de produção.
    • Leitura por caminho (por exemplo, spark.read.load(path) ou spark.read.format("delta").load(path)) devolve dados reais de produção em vez da simulação.
    • A leitura de um conector liga-se à fonte real de produção. Isto inclui o Kafka (lê dos corretores reais) e o Auto Loader (cloudFiles, que lê do caminho real de armazenamento na cloud). Nenhum deles é redirecionado para os seus dados simulados.

  • Não utilize a função event_log() com valor de tabela num teste unitário de pipeline. No modo de teste, event_log() não é redirecionado para o registo de eventos da sua execução de teste. Pode devolver o registo de eventos de produção ou o registo de eventos registado anteriormente, pelo que as asserções sobre esse registo podem aceder a dados de produção. Em vez disso, use o event_log_table_name devolvido pela execução e consulte-o através de test_spark. event_log_table_name pode ser None (por exemplo, se o nome da tabela do registo de eventos não puder ser resolvido), por isso verifique antes de consultar:

    status = test_pipeline.run(test_spark, set(["catalog.schema.table"]))
assert status.event_log_table_name is not None
events = test_spark.table(status.event_log_table_name)

    Não afirme status.is_success antes de ler o registo de eventos se o seu objetivo é diagnosticar uma atualização falhada. O registo de eventos é frequentemente o que inspecionas para perceber porque é que uma atualização falhou.

Governação e operações da DDL

  • Mutações de catálogo, esquema, permissões, propriedade, etiquetas e políticas não são suportadas. Isto inclui CREATE/DROP/ALTER CATALOG,CREATE/DROP/ALTER SCHEMA (incluindo SET MANAGED LOCATION), GRANT/REVOKE, ALTER ... OWNER TO,SET/UNSET TAGS , , e .CREATE/DROP POLICY Algumas instruções SQL executadas através de test_spark são rejeitadas por uma questão de defesa em profundidade; outras formas, ou as mesmas operações invocadas através de APIs diretas, podem aceder a objetos reais de produção. Não confie nestes guardas como limite de isolamento. Mantenha estas instruções fora do seu código de teste e de qualquer código de pipeline executado pelas saídas selecionadas.

Limitações operacionais

  • A execução simultânea não é suportada: Executar um teste e uma atualização do pipeline ao mesmo tempo não é suportado, e o sistema não o impede. Não há coordenação entre os dois, por isso executá-los em simultâneo pode disputar recursos, degradando gravemente o desempenho da atualização de produção ou fazendo com que o teste não comece. Não inicie um teste enquanto o pipeline está a correr uma atualização (ou inicie uma atualização enquanto o teste está a decorrer); Espere que qualquer atualização em curso termine antes de executar os testes.
  • Esquemas temporários após terminação anormal: Cada execução de teste cria um esquema temporário (denominado redirecting_<id>) no catálogo predefinido do pipeline e descarta-o automaticamente quando a execução termina. Se uma execução terminar de forma anómala (por exemplo, se os recursos de computação falharem a meio da execução), o esquema temporário pode ficar para trás, contendo as tabelas de simulação e de saída da execução. Não afeta os dados de produção. Para recuperar armazenamento, elimine manualmente quaisquer esquemas remanescentes cujos nomes comecem redirecting_ no catálogo padrão do pipeline.
  • As execuções de teste consomem computação: As execuções de teste são executadas no cálculo do pipeline e são faturadas como atualizações normais do pipeline. Não existe medição separada para os testes.
  • A atualização completa não é suportada: Apenas a atualização seletiva está disponível. test_pipeline.run() atualiza os outputs que selecionar (ou todos os outputs quando não for passada qualquer seleção); a atualização completa e a seleção full-refresh não estão implementadas.

Limitações de autoria e fidelidade

  • Execução apenas com o editor: Os testes devem ser executados a partir do Lakeflow Pipelines Editor baseado na web.
  • Apenas testes em Python: Os testes devem ser escritos em Python. Podes testar pipelines SQL, mas os testes em si têm de estar escritos em Python.
  • Fidelidade da governação: Os dados simulados não herdam filtros de linhas nem máscaras de colunas definidas nas tabelas de produção que vêm substituir. Os resultados dos testes refletem os dados simulados exatamente como os fornece e podem diferir da forma como a mesma consulta se comporta em dados de produção sujeitos a governação.

Passo 1: Atualizar as definições do pipeline

Configure o pipeline para correr no canal PREVIEW em modo desencadeado.

  1. Na interface, abra o seu pipeline e clique em Definições>Definições avançadas>Canal>Pré-visualização
  2. Defina o modo de pipeline como Acionado (não utilize Contínuo).

Alternativamente, edita diretamente as definições do pipeline em JSON:

"continuous": false,
"channel": "PREVIEW"

Passo 2: Criar um ficheiro de teste

No Editor de Pipelines Lakeflow, clique no + botão (adicionar) e selecione Testar. Isto cria um ficheiro de teste (e a tests pasta, caso ainda não exista) que não está incluído no código-fonte do pipeline. Não precisas de criar a tests pasta tu próprio.

Adicionar o menu de ativos do pipeline que mostra a opção Testar para criar um ficheiro pytest.

Passo 3: Gerar testes

O Genie Code pode gerar andaimes de teste:

  • Dentro do ficheiro de teste, clique no botão Gerar testes .

    Esvaziar ficheiro de teste com o botão Gerar testes.

  • Em alternativa, usa /tests no modo de agente do Genie Code.

    ficheiro de teste preenchido por Genie Code com testes unitários baseados em TestPipeline.

Use Genie Code para gerar código-base e depois adapte-o aos seus casos limite.

Em alternativa, podes escrever o código do teste tu próprio. Adicione as seguintes importações no topo de cada ficheiro de teste:

import pytest
from pyspark.pipelines.testing import TestPipeline, test_spark

test_pipeline = TestPipeline.active()

Passo 4: Fazer testes

Execute testes a partir do Editor Lakeflow Pipelines:

  • Clique no botão ícone Reproduzir. (play), na margem ao lado de uma função de teste, para executar um teste individual.
  • Clique em executar testes no ficheiro no topo do ficheiro de teste para executar todos os testes nesse ficheiro.

Os resultados dos testes (sucesso ou fracasso) aparecem no painel inferior do Editor. Analise os erros de asserção para depurar falhas.

Testar APIs

API Description
TestPipeline.active() Devolve um objeto TestPipeline correspondente ao pipeline que está atualmente a ser editado no Editor de Pipelines do Lakeflow. Este objeto é uma referência ao pipeline, incluindo o seu código-fonte, configurações, catálogo/esquema padrão, etc.
test_pipeline.run(test_spark, set([table_names])) Executa sincronizadamente uma atualização do pipeline, realizando uma atualização seletiva se os nomes das tabelas forem especificados. Retorna após a execução do pipeline ser bem-sucedida ou terminar com uma exceção.
test_spark fixação Cria uma SparkSession de teste com redirecionamento de tabelas do catálogo que redireciona automaticamente operações de leitura e escrita que façam referência a uma tabela pelo nome (por exemplo, spark.read.table("catalog.schema.table") ou df.write.saveAsTable("catalog.schema.table")) para um esquema de teste temporário. O redirecionamento aplica-se apenas a operações de tabela baseadas em nomes; Não cobre leituras ou escritas endereçadas por caminho ou através de um conector, que atuam diretamente sobre o sistema real. Consulte Limitações.

Criar dados simulados

Pode simular dados de entrada usando SQL ou createDataFrame:

# Option 1: Using SQL
test_spark.sql("""
    CREATE TABLE catalog.schema.table_name AS
    SELECT * FROM VALUES
        (1, 'value1'),
        (2, 'value2')
    AS t(id, name)
""")

# Option 2: Using createDataFrame
df = test_spark.createDataFrame(
    [(1, 'value1'), (2, 'value2')],
    schema=["id", "name"]
)
df.write.saveAsTable("catalog.schema.table_name")

Para gerar volumes maiores de dados sintéticos realistas, pode usar a biblioteca Faker . Execute %pip install faker primeiro no seu pipeline e depois crie um DataFrame a partir de UDFs suportadas pelo Faker:

# Option 3: Using Faker for synthetic data
from pyspark.sql import functions as F
from faker import Faker

fake = Faker()
fake_firstname = F.udf(fake.first_name)
fake_lastname = F.udf(fake.last_name)
fake_email = F.udf(fake.ascii_company_email)

df = (
    test_spark.range(0, 100)
    .withColumn("firstname", fake_firstname())
    .withColumn("lastname", fake_lastname())
    .withColumn("email", fake_email())
)
df.write.saveAsTable("catalog.schema.table_name")

Executar o pipeline ou tabelas específicas

# Run specific tables
test_pipeline.run(test_spark, set(["catalog.schema.table1", "catalog.schema.table2"]))

# Run all tables in the pipeline
test_pipeline.run(test_spark)

Exemplos

Exemplo 1: Testar agregações com contagem de linhas, esquema e tratamento de nulos

Objetivo: Validar corretamente a agregação de utilizadores, contar os utilizadores pelo tipo, tratar de emails nulos e produzir o esquema esperado.

Transformações de canalização:

Estas transformações criam um pipeline simples de duas tabelas: users seleciona os dados dos utilizadores, counts agrupa os utilizadores por tipo e conta o total de utilizadores e emails válidos.

from pyspark import pipelines as dp
from pyspark.sql.functions import col, count, count_if

@dp.table
def users():
    return (
        spark.read.table("catalog.schema.wanderbricks_users")
        .select("user_id", "email", "name", "user_type")
    )

@dp.table
def counts():
    return (
        spark.read.table("catalog.schema.users")
        .withColumn("valid_email", col("email").isNotNull())
        .groupBy("user_type")
        .agg(
            count("user_id").alias("total_count"),
            count_if("valid_email").alias("count_valid_emails")
        )
    )

Testes:

Estes testes validam a contagem de linhas, a estrutura do esquema, o tratamento de nulos e a lógica de agregação, criando dados simulados do utilizador com nulos intencionais e executando o pipeline isoladamente.

import pytest
from pyspark.pipelines.testing import TestPipeline, test_spark
from pyspark.testing import assertDataFrameEqual

test_pipeline = TestPipeline.active()

# Mock data fixture
def mock_users(session):
    session.sql("""
        CREATE TABLE catalog.schema.wanderbricks_users AS
        SELECT * FROM VALUES
            (1, 'alice@example.com', 'Alice', 'admin'),
            (2, NULL, 'Bob', 'user'),
            (3, 'charlie@example.com', 'Charlie', 'user'),
            (4, NULL, 'Dana', 'admin')
        AS t(user_id, email, name, user_type)
    """)

# Test 1: Row count
def test_users_row_count(test_spark):
    mock_users(test_spark)
    test_pipeline.run(test_spark, set(["catalog.schema.users"]))
    result = test_spark.table("catalog.schema.users")
    assert result.count() == 4

# Test 2: Schema validation
def test_users_schema(test_spark):
    mock_users(test_spark)
    test_pipeline.run(test_spark, set(["catalog.schema.users"]))
    result = test_spark.table("catalog.schema.users")
    expected_fields = {"user_id", "email", "name", "user_type"}
    actual_fields = set(f.name for f in result.schema.fields)
    assert expected_fields == actual_fields

# Test 3: Null handling
def test_users_null_handling(test_spark):
    mock_users(test_spark)
    test_pipeline.run(test_spark, set(["catalog.schema.users"]))
    result = test_spark.table("catalog.schema.users")
    null_emails = result.filter("email IS NULL").count()
    assert null_emails == 2

# Test 4: Aggregation
def test_counts(test_spark):
    mock_users(test_spark)
    # Run both tables since counts depends on users
    test_pipeline.run(test_spark, set(["catalog.schema.users", "catalog.schema.counts"]))
    result = test_spark.table("catalog.schema.counts")
    # Check counts for each user_type
    admin_row = result.filter("user_type = 'admin'").collect()[0]
    user_row = result.filter("user_type = 'user'").collect()[0]
    assert admin_row["total_count"] == 2
    assert admin_row["count_valid_emails"] == 1
    assert user_row["total_count"] == 2
    assert user_row["count_valid_emails"] == 1

# Test 5: Full DataFrame comparison with assertDataFrameEqual
def test_counts_full_dataframe(test_spark):
    mock_users(test_spark)
    test_pipeline.run(test_spark, set(["catalog.schema.users", "catalog.schema.counts"]))
    result = test_spark.table("catalog.schema.counts")
    expected = test_spark.createDataFrame(
        [("admin", 2, 1), ("user", 2, 1)],
        schema=["user_type", "total_count", "count_valid_emails"]
    )
    assertDataFrameEqual(result, expected)

Exemplo 2: Testar o Auto CDC

Objetivo: Validar que o Auto CDC processa corretamente o feed de alterações com inserções e atualizações.

Transformação do pipeline:

Esta transformação configura o Auto CDC com base num feed de alterações, que lê alterações em fluxo e as aplica à tabela de destino como SCD de Tipo 1 (mantém apenas a versão mais recente).

from pyspark import pipelines as dp
from pyspark.sql.functions import col

@dp.view
def users():
    return spark.readStream.table("catalog.schema.change_feed")

dp.create_streaming_table("target_autocdc")
dp.create_auto_cdc_flow(
    target="target_autocdc",
    source="users",
    keys=["userId"],
    sequence_by=col("ts"),
    stored_as_scd_type=1
)

Testes:

O primeiro teste cria um feed de alterações simulado com múltiplos registos para o mesmo userId (simulando uma atualização) e verifica se apenas o registo mais recente é mantido no destino. O segundo teste simula eventos que chegam tardiamente e fora de ordem, executando o pipeline, acrescentando mais eventos ao feed de alterações e executando o pipeline novamente.

import pytest
from pyspark.pipelines.testing import TestPipeline, test_spark

test_pipeline = TestPipeline.active()

# Test 1: Standard inserts and updates
def test_auto_cdc_flow(test_spark):
    # Create a mock change feed table
    test_spark.sql("""
        CREATE TABLE catalog.schema.change_feed AS
        SELECT * FROM VALUES
            (1, 'Alice', 1000),
            (2, 'Bob', 1001),
            (1, 'Alice Updated', 1002)
        AS t(userId, name, ts)
    """)
    # Run the pipeline
    test_pipeline.run(test_spark, set(["catalog.schema.target_autocdc"]))
    # Read the output
    result = test_spark.table("catalog.schema.target_autocdc")
    # Verify two users exist
    user_ids = set(row["userId"] for row in result.collect())
    assert user_ids == {1, 2}
    # Verify latest record for userId=1 has ts=1002
    latest_user1 = result.filter("userId = 1").collect()[0]
    assert latest_user1["ts"] == 1002
    assert latest_user1["name"] == "Alice Updated"
    # Verify userId=2 has ts=1001
    user2 = result.filter("userId = 2").collect()[0]
    assert user2["ts"] == 1001

# Test 2: Late-arriving and out-of-order events
def test_auto_cdc_late_arriving(test_spark):
    # First batch of change events
    test_spark.sql("""
        CREATE TABLE catalog.schema.change_feed AS
        SELECT * FROM VALUES
            (1, 'Alice', 1000),
            (2, 'Bob', 1001)
        AS t(userId, name, ts)
    """)
    # Run the pipeline with the initial batch
    test_pipeline.run(test_spark, set(["catalog.schema.target_autocdc"]))

    # Append late-arriving events to the change feed:
    # - A newer event for userId=1 (ts=1003) that arrived after the first run
    # - A stale event for userId=2 (ts=999) with a timestamp older than what is already applied
    test_spark.sql("""
        INSERT INTO catalog.schema.change_feed VALUES
            (1, 'Alice Updated', 1003),
            (2, 'Bob (stale)', 999)
    """)
    # Re-run the pipeline. sequence_by=ts ensures stale events do not overwrite newer state.
    test_pipeline.run(test_spark, set(["catalog.schema.target_autocdc"]))

    result = test_spark.table("catalog.schema.target_autocdc")
    # userId=1 should reflect the newer late-arriving event
    alice = result.filter("userId = 1").collect()[0]
    assert alice["ts"] == 1003
    assert alice["name"] == "Alice Updated"
    # userId=2 should be unchanged: the stale event with an older ts is ignored
    bob = result.filter("userId = 2").collect()[0]
    assert bob["ts"] == 1001
    assert bob["name"] == "Bob"

Exemplo 3: Testar o Auto CDC a partir de snapshot

Objetivo: Validar que o CDC processa corretamente as alterações de snapshots, incluindo inserções, atualizações e eliminações.

Transformação do pipeline:

Esta transformação cria o Auto CDC a partir do snapshot, que lê de uma tabela de snapshots e acompanha as alterações ao longo do tempo como SCD Tipo 2 (mantém o histórico completo).

from pyspark import pipelines as dp

@dp.view(name="source")
def source():
    return spark.read.table("catalog.schema.snapshot")

dp.create_streaming_table("catalog.schema.target")
dp.create_auto_cdc_from_snapshot_flow(
    target="target",
    source="source",
    keys=["userId"],
    stored_as_scd_type=2
)

Teste:

Este teste cria um snapshot inicial, executa o pipeline e depois simula uma atualização do snapshot, truncando e inserindo novos dados, para verificar se o CDC captura todas as alterações.

import pytest
from pyspark.pipelines.testing import TestPipeline, test_spark

test_pipeline = TestPipeline.active()

def test_auto_cdc_from_snapshot_flow(test_spark):
    # Create initial snapshot
    test_spark.sql("""
        CREATE TABLE catalog.schema.snapshot AS
        SELECT * FROM VALUES
            (1, 'Alice', '2024-01-01'),
            (2, 'Bob', '2024-01-02')
        AS t(userId, name, created_at)
    """)
    # Run the pipeline
    test_pipeline.run(test_spark, set(["catalog.schema.target"]))
    # Simulate a new snapshot by truncating and inserting updated data
    test_spark.sql("TRUNCATE TABLE catalog.schema.snapshot")
    test_spark.sql("INSERT INTO catalog.schema.snapshot VALUES (2, 'Bob', '2024-01-03')")
    test_pipeline.run(test_spark, set(["catalog.schema.target"]))
    # Verify SCD Type 2: should have 3 rows (original Alice, original Bob, updated Bob)
    result = test_spark.table("catalog.schema.target")
    assert result.count() == 3
    user_ids = [row["userId"] for row in result.collect()]
    assert set(user_ids) == {1, 2}

Exemplo 4: Teste de junções e expectativas

Objetivo: Validar que as junções funcionam corretamente e que as expectativas filtram dados inválidos.

Transformação do pipeline:

Esta transformação combina as imagens da propriedade com as comodidades e aplica um critério para filtrar as imagens carregadas antes de janeiro de 2024.

from pyspark import pipelines as dp

@dp.table
@dp.expect_or_drop("uploaded after Jan 2024", "uploaded_at > '2024-01-01'")
def property_images_amenities_join():
    return (
        spark.read.table("catalog.schema.property_images")
        .join(
            spark.read.table("catalog.schema.property_amenities"),
            on="property_id",
            how="inner"
        )
    )

Testes:

Estes testes verificam que a junção produz o número correto de linhas e que a expectativa filtra com sucesso registos com datas de upload inválidas.

import pytest
from pyspark.pipelines.testing import TestPipeline, test_spark

test_pipeline = TestPipeline.active()

# Mock property datasets
def mock_properties(session):
    session.sql("""
        CREATE TABLE catalog.schema.property_images AS
        SELECT * FROM VALUES
            (101, 'img1.jpg', '2024-02-01'),
            (102, 'img2.jpg', '2024-01-15'),
            (103, 'img3.jpg', '2024-12-20')
        AS t(property_id, image_url, uploaded_at)
    """)
    session.sql("""
        CREATE TABLE catalog.schema.property_amenities AS
        SELECT * FROM VALUES
            (101, 'wifi'),
            (102, 'pool'),
            (103, 'parking')
        AS t(property_id, amenity)
    """)

# Test 1: Join
def test_property_join(test_spark):
    mock_properties(test_spark)
    test_pipeline.run(test_spark, set(["catalog.schema.property_images_amenities_join"]))
    result = test_spark.table("catalog.schema.property_images_amenities_join")
    # Should have 3 rows after join
    assert result.count() == 3
    # Check all property_ids are present
    property_ids = set(row["property_id"] for row in result.collect())
    assert property_ids == {101, 102, 103}

# Test 2: Expectation
def test_property_expectation(test_spark):
    mock_properties(test_spark)
    # Add a row with uploaded_at before Jan 2024
    test_spark.sql("""
        INSERT INTO catalog.schema.property_images VALUES (104, 'img4.jpg', '2023-12-31')
    """)
    # Add a matching row in the amenities table for the join
    test_spark.sql("""
        INSERT INTO catalog.schema.property_amenities VALUES (104, 'gym')
    """)
    test_pipeline.run(test_spark, set(["catalog.schema.property_images_amenities_join"]))
    result = test_spark.table("catalog.schema.property_images_amenities_join")
    # Only property_ids with uploaded_at > '2024-01-01' should be present
    valid_ids = set(row["property_id"] for row in result.collect())
    assert 104 not in valid_ids
    assert valid_ids == {101, 102, 103}
  • Last updated on