@@ -32,8 +32,7 @@ describe('clickhouseClient', () => {
3232 {
3333 inferenceId : 'abc-123' ,
3434 episodeId : 'episode-123' ,
35- // pullRequestId: 42
36- // I think we're missing something here definitely
35+ pullRequestId : 42
3736 } ,
3837 defaultConfig ,
3938 { client }
@@ -43,125 +42,34 @@ describe('clickhouseClient', () => {
4342 table : 'tensorzero.inference_records' ,
4443 values : [
4544 {
45+ inference_id : 'abc-123' ,
4646 episode_id : 'episode-123' ,
47- pull_request_id : 42 ,
48- inference_id : 'abc-123'
47+ pull_request_id : 42
4948 }
5049 ] ,
5150 format : 'JSONEachRow'
5251 } )
53- expect ( client . close ) . not . toHaveBeenCalled ( )
5452 } )
5553
56- it ( 'queries inference records with parameter binding' , async ( ) => {
57- const expectedRecords = [
58- {
59- inference_id : 'xyz' ,
60- pull_request_id : 77 ,
61- created_at : '2024-01-01T00:00:00Z' ,
62- original_pull_request_url : 'https://github.com/org/repo/pull/77'
63- }
54+ it ( 'queries episode records for a pull request' , async ( ) => {
55+ const mockRows = [
56+ { episode_id : 'ep1' , pull_request_id : 123 } ,
57+ { episode_id : 'ep2' , pull_request_id : 123 }
6458 ]
65- const jsonMock = jest
66- . fn < ( ) => Promise < unknown > > ( )
67- . mockResolvedValue ( expectedRecords )
68- // @ts -expect-error(Mock type is inaccurate)
69- client . query . mockResolvedValueOnce ( { json : jsonMock } )
70-
71- const records = await getPullRequestToEpisodeRecords ( 77 , defaultConfig , {
72- client
73- } )
74-
75- expect ( client . query ) . toHaveBeenCalledWith ( {
76- query :
77- 'SELECT episode_id, pull_request_id, created_at, original_pull_request_url FROM tensorzero.inference_records WHERE pull_request_id = {pullRequestId:UInt64}' ,
78- query_params : { pullRequestId : 77 } ,
79- format : 'JSONEachRow'
80- } )
81- expect ( jsonMock ) . toHaveBeenCalledTimes ( 1 )
82- expect ( records ) . toEqual ( expectedRecords )
83- } )
84-
85- it ( 'throws when the table name fails validation' , async ( ) => {
86- await expect (
87- createPullRequestToInferenceRecord (
88- {
89- inferenceId : 'abc' ,
90- episodeId : 'episode-123' ,
91- pullRequestId : 1
92- } ,
93- { ...defaultConfig , table : 'invalid-table!' }
94- )
95- ) . rejects . toThrow ( 'ClickHouse table name must contain only' )
96- } )
97-
98- it ( 'validates missing URL' , async ( ) => {
99- await expect (
100- createPullRequestToInferenceRecord (
101- {
102- inferenceId : 'abc' ,
103- episodeId : 'episode-123' ,
104- pullRequestId : 1
105- } ,
106- { ...defaultConfig , url : ' ' }
107- )
108- ) . rejects . toThrow ( 'ClickHouse URL is required' )
109- } )
11059
111- it ( 'propagates insert failures from the injected client without closing it' , async ( ) => {
112- client . insert . mockRejectedValueOnce ( new Error ( 'insert failed' ) )
113-
114- await expect (
115- createPullRequestToInferenceRecord (
116- {
117- inferenceId : 'abc' ,
118- episodeId : 'episode-123' ,
119- pullRequestId : 1
120- } ,
121- defaultConfig ,
122- { client }
123- )
124- ) . rejects . toThrow ( 'insert failed' )
60+ client . query . mockResolvedValue ( {
61+ json : jest . fn ( ) . mockResolvedValue ( mockRows )
62+ } as never )
12563
126- expect ( client . close ) . not . toHaveBeenCalled ( )
127- } )
128-
129- it ( 'queries episode records with parameter binding' , async ( ) => {
130- const expectedRecords = [
131- {
132- episode_id : 'episode-789' ,
133- pull_request_id : 55 ,
134- created_at : '2024-02-01T00:00:00Z' ,
135- original_pull_request_url : 'https://github.com/org/repo/pull/55'
136- }
137- ]
138- const jsonMock = jest
139- . fn < ( ) => Promise < unknown > > ( )
140- . mockResolvedValue ( expectedRecords )
141- // @ts -expect-error(Mock type is inaccurate)
142- client . query . mockResolvedValueOnce ( { json : jsonMock } )
143-
144- const records = await getPullRequestToEpisodeRecords ( 55 , defaultConfig , {
64+ const result = await getPullRequestToEpisodeRecords ( 123 , defaultConfig , {
14565 client
14666 } )
14767
68+ expect ( result ) . toEqual ( mockRows )
14869 expect ( client . query ) . toHaveBeenCalledWith ( {
149- query :
150- 'SELECT episode_id, pull_request_id, created_at, original_pull_request_url FROM tensorzero.inference_records WHERE pull_request_id = {pullRequestId:UInt64}' ,
151- query_params : { pullRequestId : 55 } ,
70+ query : expect . stringContaining ( 'SELECT episode_id, pull_request_id' ) ,
71+ query_params : { pullRequestId : 123 } ,
15272 format : 'JSONEachRow'
15373 } )
154- expect ( jsonMock ) . toHaveBeenCalledTimes ( 1 )
155- expect ( records ) . toEqual ( expectedRecords )
156- } )
157-
158- it ( 'propagates query failures for episode records without closing the client' , async ( ) => {
159- client . query . mockRejectedValueOnce ( new Error ( 'query failed' ) )
160-
161- await expect (
162- getPullRequestToEpisodeRecords ( 123 , defaultConfig , { client } )
163- ) . rejects . toThrow ( 'query failed' )
164-
165- expect ( client . close ) . not . toHaveBeenCalled ( )
16674 } )
16775} )
0 commit comments