Coverage for tests \ dataset_test.py: 100%

1import json 

2from os.path import isfile 

3from tempfile import TemporaryDirectory 

4 

5from pandas import DataFrame 

6from pyarrow.dataset import dataset 

7 

8from gtrends_collection import write_to_dataset 

9 

10 

11def test_dataset(): 

12 with TemporaryDirectory() as tempdir: 

13 data_dir = tempdir + "/data" 

14 data = DataFrame( 

15 { 

16 "term": ["a", "b"], 

17 "value": [0] * 2, 

18 "date": ["2004-01-01"] * 2, 

19 "location": ["US"] * 2, 

20 "retrieved": ["2025-03-07"] * 2, 

21 } 

22 ) 

23 write_to_dataset(data, data_dir) 

24 assert dataset(data_dir).scanner(["term"]).to_table()["term"].to_pylist() == data["term"].to_list() 

25 

26 log_file = tempdir + "/status.json" 

27 assert isfile(log_file) 

28 with open(log_file, "r", encoding="utf-8") as file: 

29 state = json.load(file) 

30 assert data_dir.replace("\\", "/") + "/term=b/part-0.parquet" in state