Skip to content

Commit 91157fb

Browse files
fix mock execute result in test
1 parent b5e35fa commit 91157fb

File tree

1 file changed

+23
-8
lines changed

1 file changed

+23
-8
lines changed

tests/system/small/test_anywidget.py

Lines changed: 23 additions & 8 deletions
Original file line numberDiff line numberDiff 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

131146
def _assert_html_matches_pandas_slice(
132147
table_html: str,

0 commit comments

Comments
 (0)