1+ import pytest
2+ from unittest .mock import MagicMock , Mock , patch
3+ from pyiceberg .table import DataScan
4+ from pyiceberg .expressions import AlwaysTrue
5+
6+ class DummyFile :
7+ def __init__ (self , record_count ):
8+ self .record_count = record_count
9+
10+ class DummyTask :
11+ def __init__ (self , record_count , residual = None , delete_files = None ):
12+ self .file = DummyFile (record_count )
13+ self .residual = residual if residual is not None else AlwaysTrue ()
14+ self .delete_files = delete_files or []
15+
16+ def test_count_basic ():
17+ # Create a mock table with the necessary attributes
18+ table = Mock (spec = DataScan )
19+
20+ # Mock the plan_files method to return our dummy task
21+ task = DummyTask (42 , residual = AlwaysTrue (), delete_files = [])
22+ table .plan_files = MagicMock (return_value = [task ])
23+
24+ # Import and call the actual count method
25+ from pyiceberg .table import DataScan as ActualDataScan
26+ table .count = ActualDataScan .count .__get__ (table , ActualDataScan )
27+
28+ assert table .count () == 42
29+
30+ def test_count_empty ():
31+ # Create a mock table with the necessary attributes
32+ table = Mock (spec = DataScan )
33+
34+ # Mock the plan_files method to return no tasks
35+ table .plan_files = MagicMock (return_value = [])
36+
37+ # Import and call the actual count method
38+ from pyiceberg .table import DataScan as ActualDataScan
39+ table .count = ActualDataScan .count .__get__ (table , ActualDataScan )
40+
41+ assert table .count () == 0
42+
43+ def test_count_large ():
44+ # Create a mock table with the necessary attributes
45+ table = Mock (spec = DataScan )
46+
47+ # Mock the plan_files method to return multiple tasks
48+ tasks = [
49+ DummyTask (500000 , residual = AlwaysTrue (), delete_files = []),
50+ DummyTask (500000 , residual = AlwaysTrue (), delete_files = []),
51+ ]
52+ table .plan_files = MagicMock (return_value = tasks )
53+
54+ # Import and call the actual count method
55+ from pyiceberg .table import DataScan as ActualDataScan
56+ table .count = ActualDataScan .count .__get__ (table , ActualDataScan )
57+
58+ assert table .count () == 1000000
0 commit comments