14 lines
483 B
Python
14 lines
483 B
Python
import pytest
|
|
from src.data import data_collection
|
|
|
|
def test_collect_data():
|
|
# Test the collect_data function
|
|
data = data_collection.collect_data('path_to_test_data')
|
|
|
|
# Check that the data has the expected shape
|
|
assert data.shape == (expected_number_of_rows, expected_number_of_columns)
|
|
|
|
# Check that the data has the expected columns
|
|
expected_columns = ['column1', 'column2', 'column3']
|
|
assert all(column in data.columns for column in expected_columns)
|