diff --git a/sqlmesh/core/test/definition.py b/sqlmesh/core/test/definition.py index a12dafec19..4336e00ce0 100644 --- a/sqlmesh/core/test/definition.py +++ b/sqlmesh/core/test/definition.py @@ -644,16 +644,16 @@ def _create_df( return self._execute(query) rows = values["rows"] + columns_str: t.Optional[t.List[str]] = None if columns: + columns_str = [str(c) for c in columns] referenced_columns = list(dict.fromkeys(col for row in rows for col in row)) _raise_if_unexpected_columns(columns, referenced_columns) if partial: - columns = referenced_columns + columns_str = [c for c in columns_str if c in referenced_columns] - return pd.DataFrame.from_records( - rows, columns=[str(c) for c in columns] if columns else None - ) + return pd.DataFrame.from_records(rows, columns=columns_str) def _add_missing_columns( self, query: exp.Query, all_columns: t.Optional[t.Collection[str]] = None diff --git a/tests/core/test_test.py b/tests/core/test_test.py index d889c7bb33..56a44cc955 100644 --- a/tests/core/test_test.py +++ b/tests/core/test_test.py @@ -770,6 +770,42 @@ def test_partial_data_column_order(sushi_context: Context) -> None: ).run() ) + # - output df must differ if sorted by (id, event_date) vs. (event_date, id) + # - output partial must be true + _check_successful_or_raise( + _create_test( + body=load_yaml( + """ +test_foo: + model: sushi.foo + inputs: + sushi.items: + - id: 9876 + event_date: 2020-01-01 + - id: 1234 + name: hello + event_date: 2020-01-02 + outputs: + partial: true + query: + - event_date: 2020-01-01 + id: 9876 + - event_date: 2020-01-02 + id: 1234 + name: hello + """ + ), + test_name="test_foo", + model=sushi_context.upsert_model( + _create_model( + "SELECT id, name, price, event_date FROM sushi.items", + default_catalog=sushi_context.default_catalog, + ) + ), + context=sushi_context, + ).run() + ) + def test_partial_data_missing_schemas(sushi_context: Context) -> None: _check_successful_or_raise(