File tree Expand file tree Collapse file tree 1 file changed +23
-8
lines changed
Expand file tree Collapse file tree 1 file changed +23
-8
lines changed Original file line number Diff line number Diff line change @@ -117,16 +117,31 @@ def mock_execute_result_with_params(
117117 """
118118 Mocks an execution result with configurable total_rows and arrow_batches.
119119 """
120- from bigframes .session .executor import ExecuteResult
121-
122- return ExecuteResult (
123- iter (arrow_batches_val ),
124- schema = schema ,
125- query_job = None ,
126- total_bytes = None ,
127- total_rows = total_rows_val ,
120+ from bigframes .session .executor import (
121+ ExecuteResult ,
122+ ExecutionMetadata ,
123+ ResultsIterator ,
128124 )
129125
126+ class MockExecuteResult (ExecuteResult ):
127+ @property
128+ def execution_metadata (self ) -> ExecutionMetadata :
129+ return ExecutionMetadata ()
130+
131+ @property
132+ def schema (self ):
133+ return schema
134+
135+ def batches (self ) -> ResultsIterator :
136+ return ResultsIterator (
137+ arrow_batches_val ,
138+ self .schema ,
139+ total_rows_val ,
140+ None ,
141+ )
142+
143+ return MockExecuteResult ()
144+
130145
131146def _assert_html_matches_pandas_slice (
132147 table_html : str ,
You can’t perform that action at this time.
0 commit comments