Skip to content

Commit 1a9b7de

Browse files
jankucajankuca
committed
test: add tests for updated sql env var provider
1 parent 94ba87e commit 1a9b7de

File tree

1 file changed

+385
-0
lines changed

1 file changed

+385
-0
lines changed

‎src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts‎

Lines changed: 385 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,385 @@
1+
import { assert } from 'chai';
2+
import { instance, mock, when } from 'ts-mockito';
3+
import { CancellationTokenSource, EventEmitter, NotebookDocument, Uri } from 'vscode';
4+
5+
import { IDisposableRegistry } from '../../common/types';
6+
import { SqlIntegrationEnvironmentVariablesProvider } from './sqlIntegrationEnvironmentVariablesProvider';
7+
import { IIntegrationStorage, IPlatformDeepnoteNotebookManager, IPlatformNotebookEditorProvider } from './types';
8+
import { DATAFRAME_SQL_INTEGRATION_ID } from './integrationTypes';
9+
import { DatabaseIntegrationConfig } from '@deepnote/database-integrations';
10+
import type { DeepnoteProject } from '../../deepnote/deepnoteTypes';
11+
12+
/**
13+
* Helper function to create a minimal DeepnoteProject for testing
14+
*/
15+
function createMockProject(
16+
projectId: string,
17+
integrations: Array<{ id: string; name: string; type: string }> = []
18+
): DeepnoteProject {
19+
return {
20+
metadata: {
21+
createdAt: '2023-01-01T00:00:00Z',
22+
modifiedAt: '2023-01-02T00:00:00Z'
23+
},
24+
project: {
25+
id: projectId,
26+
name: 'Test Project',
27+
notebooks: [],
28+
integrations
29+
},
30+
version: '1.0'
31+
};
32+
}
33+
34+
suite('SqlIntegrationEnvironmentVariablesProvider', () => {
35+
let provider: SqlIntegrationEnvironmentVariablesProvider;
36+
let integrationStorage: IIntegrationStorage;
37+
let notebookEditorProvider: IPlatformNotebookEditorProvider;
38+
let notebookManager: IPlatformDeepnoteNotebookManager;
39+
let disposables: IDisposableRegistry;
40+
let onDidChangeIntegrationsEmitter: EventEmitter<void>;
41+
42+
setup(() => {
43+
integrationStorage = mock<IIntegrationStorage>();
44+
notebookEditorProvider = mock<IPlatformNotebookEditorProvider>();
45+
notebookManager = mock<IPlatformDeepnoteNotebookManager>();
46+
disposables = [];
47+
48+
onDidChangeIntegrationsEmitter = new EventEmitter<void>();
49+
when(integrationStorage.onDidChangeIntegrations).thenReturn(onDidChangeIntegrationsEmitter.event);
50+
51+
provider = new SqlIntegrationEnvironmentVariablesProvider(
52+
instance(integrationStorage),
53+
instance(notebookEditorProvider),
54+
instance(notebookManager),
55+
disposables
56+
);
57+
});
58+
59+
teardown(() => {
60+
disposables.forEach((d) => d.dispose());
61+
onDidChangeIntegrationsEmitter.dispose();
62+
});
63+
64+
suite('getEnvironmentVariables', () => {
65+
test('Returns empty object when resource is undefined', async () => {
66+
const result = await provider.getEnvironmentVariables(undefined);
67+
68+
assert.deepStrictEqual(result, {});
69+
});
70+
71+
test('Returns empty object when cancellation token is already cancelled', async () => {
72+
const tokenSource = new CancellationTokenSource();
73+
tokenSource.cancel();
74+
const resource = Uri.file('/test/notebook.deepnote');
75+
76+
const result = await provider.getEnvironmentVariables(resource, tokenSource.token);
77+
78+
assert.deepStrictEqual(result, {});
79+
tokenSource.dispose();
80+
});
81+
82+
test('Returns empty object when no notebook is found for resource', async () => {
83+
const resource = Uri.file('/test/notebook.deepnote');
84+
when(notebookEditorProvider.findAssociatedNotebookDocument(resource)).thenReturn(undefined);
85+
86+
const result = await provider.getEnvironmentVariables(resource);
87+
88+
assert.deepStrictEqual(result, {});
89+
});
90+
91+
test('Returns empty object when notebook has no project ID in metadata', async () => {
92+
const resource = Uri.file('/test/notebook.deepnote');
93+
const notebook = mock<NotebookDocument>();
94+
when(notebook.metadata).thenReturn({});
95+
when(notebookEditorProvider.findAssociatedNotebookDocument(resource)).thenReturn(instance(notebook));
96+
97+
const result = await provider.getEnvironmentVariables(resource);
98+
99+
assert.deepStrictEqual(result, {});
100+
});
101+
102+
test('Returns empty object when project is not found in notebook manager', async () => {
103+
const resource = Uri.file('/test/notebook.deepnote');
104+
const notebook = mock<NotebookDocument>();
105+
when(notebook.metadata).thenReturn({ deepnoteProjectId: 'project-123' });
106+
when(notebookEditorProvider.findAssociatedNotebookDocument(resource)).thenReturn(instance(notebook));
107+
when(notebookManager.getOriginalProject('project-123')).thenReturn(undefined);
108+
109+
const result = await provider.getEnvironmentVariables(resource);
110+
111+
assert.deepStrictEqual(result, {});
112+
});
113+
114+
test('Returns only DuckDB integration when project has no integrations', async () => {
115+
const resource = Uri.file('/test/notebook.deepnote');
116+
const notebook = mock<NotebookDocument>();
117+
const project = createMockProject('project-123', []);
118+
119+
when(notebook.metadata).thenReturn({ deepnoteProjectId: 'project-123' });
120+
when(notebookEditorProvider.findAssociatedNotebookDocument(resource)).thenReturn(instance(notebook));
121+
when(notebookManager.getOriginalProject('project-123')).thenReturn(project);
122+
123+
const result = await provider.getEnvironmentVariables(resource);
124+
125+
// Should contain DuckDB integration env vars
126+
assert.ok(Object.keys(result).length > 0, 'Should have environment variables for DuckDB');
127+
// The actual env var name depends on the database-integrations library implementation
128+
// We verify that at least one env var was generated
129+
});
130+
131+
test('Retrieves integration configs from storage for project integrations', async () => {
132+
const resource = Uri.file('/test/notebook.deepnote');
133+
const notebook = mock<NotebookDocument>();
134+
const postgresConfig: DatabaseIntegrationConfig = {
135+
id: 'postgres-1',
136+
name: 'My Postgres DB',
137+
type: 'pgsql',
138+
metadata: {
139+
host: 'localhost',
140+
port: '5432',
141+
database: 'testdb',
142+
user: 'testuser',
143+
password: 'testpass',
144+
sslEnabled: false
145+
}
146+
};
147+
const project = createMockProject('project-123', [
148+
{ id: 'postgres-1', name: 'My Postgres DB', type: 'pgsql' }
149+
]);
150+
151+
when(notebook.metadata).thenReturn({ deepnoteProjectId: 'project-123' });
152+
when(notebookEditorProvider.findAssociatedNotebookDocument(resource)).thenReturn(instance(notebook));
153+
when(notebookManager.getOriginalProject('project-123')).thenReturn(project);
154+
when(integrationStorage.getIntegrationConfig('postgres-1')).thenResolve(postgresConfig);
155+
156+
const result = await provider.getEnvironmentVariables(resource);
157+
158+
// Should contain env vars for both Postgres and DuckDB
159+
assert.ok(Object.keys(result).length > 0, 'Should have environment variables');
160+
});
161+
162+
test('Filters out null integration configs from storage', async () => {
163+
const resource = Uri.file('/test/notebook.deepnote');
164+
const notebook = mock<NotebookDocument>();
165+
const postgresConfig: DatabaseIntegrationConfig = {
166+
id: 'postgres-1',
167+
name: 'My Postgres DB',
168+
type: 'pgsql',
169+
metadata: {
170+
host: 'localhost',
171+
port: '5432',
172+
database: 'testdb',
173+
user: 'testuser',
174+
password: 'testpass',
175+
sslEnabled: false
176+
}
177+
};
178+
const project = createMockProject('project-123', [
179+
{ id: 'postgres-1', name: 'My Postgres DB', type: 'pgsql' },
180+
{ id: 'missing-integration', name: 'Missing', type: 'pgsql' }
181+
]);
182+
183+
when(notebook.metadata).thenReturn({ deepnoteProjectId: 'project-123' });
184+
when(notebookEditorProvider.findAssociatedNotebookDocument(resource)).thenReturn(instance(notebook));
185+
when(notebookManager.getOriginalProject('project-123')).thenReturn(project);
186+
when(integrationStorage.getIntegrationConfig('postgres-1')).thenResolve(postgresConfig);
187+
when(integrationStorage.getIntegrationConfig('missing-integration')).thenResolve(undefined);
188+
189+
const result = await provider.getEnvironmentVariables(resource);
190+
191+
// Should only include postgres-1 and DuckDB, not the missing integration
192+
assert.ok(Object.keys(result).length > 0, 'Should have environment variables');
193+
});
194+
195+
test('Always includes DuckDB integration in the config list', async () => {
196+
const resource = Uri.file('/test/notebook.deepnote');
197+
const notebook = mock<NotebookDocument>();
198+
const project = createMockProject('project-123', []);
199+
200+
when(notebook.metadata).thenReturn({ deepnoteProjectId: 'project-123' });
201+
when(notebookEditorProvider.findAssociatedNotebookDocument(resource)).thenReturn(instance(notebook));
202+
when(notebookManager.getOriginalProject('project-123')).thenReturn(project);
203+
204+
const result = await provider.getEnvironmentVariables(resource);
205+
206+
// DuckDB should always be included
207+
assert.ok(Object.keys(result).length > 0, 'Should have DuckDB environment variables');
208+
});
209+
210+
test('Generates environment variables for multiple integrations', async () => {
211+
const resource = Uri.file('/test/notebook.deepnote');
212+
const notebook = mock<NotebookDocument>();
213+
const postgresConfig: DatabaseIntegrationConfig = {
214+
id: 'postgres-1',
215+
name: 'Postgres DB',
216+
type: 'pgsql',
217+
metadata: {
218+
host: 'localhost',
219+
port: '5432',
220+
database: 'testdb',
221+
user: 'testuser',
222+
password: 'testpass',
223+
sslEnabled: false
224+
}
225+
};
226+
const bigqueryConfig: DatabaseIntegrationConfig = {
227+
id: 'bigquery-1',
228+
name: 'BigQuery',
229+
type: 'big-query',
230+
metadata: {
231+
authMethod: 'service-account',
232+
service_account: '{"type":"service_account","project_id":"test"}'
233+
}
234+
};
235+
const project = createMockProject('project-123', [
236+
{ id: 'postgres-1', name: 'Postgres DB', type: 'pgsql' },
237+
{ id: 'bigquery-1', name: 'BigQuery', type: 'big-query' }
238+
]);
239+
240+
when(notebook.metadata).thenReturn({ deepnoteProjectId: 'project-123' });
241+
when(notebookEditorProvider.findAssociatedNotebookDocument(resource)).thenReturn(instance(notebook));
242+
when(notebookManager.getOriginalProject('project-123')).thenReturn(project);
243+
when(integrationStorage.getIntegrationConfig('postgres-1')).thenResolve(postgresConfig);
244+
when(integrationStorage.getIntegrationConfig('bigquery-1')).thenResolve(bigqueryConfig);
245+
246+
const result = await provider.getEnvironmentVariables(resource);
247+
248+
// Should have env vars for Postgres, BigQuery, and DuckDB
249+
assert.ok(Object.keys(result).length > 0, 'Should have environment variables for all integrations');
250+
});
251+
252+
suite('Real environment variable format checks', () => {
253+
test('PostgreSQL integration generates correct SQL_* env var format', async () => {
254+
const resource = Uri.file('/test/notebook.deepnote');
255+
const notebook = mock<NotebookDocument>();
256+
const postgresConfig: DatabaseIntegrationConfig = {
257+
id: 'my-postgres',
258+
name: 'Production DB',
259+
type: 'pgsql',
260+
metadata: {
261+
host: 'db.example.com',
262+
port: '5432',
263+
database: 'production',
264+
user: 'admin',
265+
password: 'secret123',
266+
sslEnabled: true
267+
}
268+
};
269+
const project = createMockProject('project-123', [
270+
{ id: 'my-postgres', name: 'Production DB', type: 'pgsql' }
271+
]);
272+
273+
when(notebook.metadata).thenReturn({ deepnoteProjectId: 'project-123' });
274+
when(notebookEditorProvider.findAssociatedNotebookDocument(resource)).thenReturn(instance(notebook));
275+
when(notebookManager.getOriginalProject('project-123')).thenReturn(project);
276+
when(integrationStorage.getIntegrationConfig('my-postgres')).thenResolve(postgresConfig);
277+
278+
const result = await provider.getEnvironmentVariables(resource);
279+
280+
// The database-integrations library generates env vars with SQL_ prefix
281+
// and the integration ID in uppercase with hyphens replaced by underscores
282+
const expectedEnvVarName = 'SQL_MY_POSTGRES';
283+
assert.ok(result[expectedEnvVarName], `Should have ${expectedEnvVarName} env var`);
284+
285+
// The value should be a JSON string with connection details
286+
const envVarValue = result[expectedEnvVarName];
287+
assert.ok(typeof envVarValue === 'string', 'Env var value should be a string');
288+
assert.ok(envVarValue, 'Env var value should not be undefined');
289+
290+
// Parse and verify the structure
291+
const parsed = JSON.parse(envVarValue!);
292+
assert.ok(parsed.url, 'Should have url field');
293+
assert.ok(parsed.url.includes('postgresql://'), 'URL should be PostgreSQL connection string');
294+
assert.ok(parsed.url.includes('db.example.com'), 'URL should contain host');
295+
assert.ok(parsed.url.includes('5432'), 'URL should contain port');
296+
assert.ok(parsed.url.includes('production'), 'URL should contain database name');
297+
});
298+
299+
test('BigQuery integration generates correct SQL_* env var format', async () => {
300+
const resource = Uri.file('/test/notebook.deepnote');
301+
const notebook = mock<NotebookDocument>();
302+
const serviceAccountJson = JSON.stringify({
303+
type: 'service_account',
304+
project_id: 'my-gcp-project',
305+
private_key_id: 'key123',
306+
private_key: '-----BEGIN PRIVATE KEY-----\ntest\n-----END PRIVATE KEY-----\n',
307+
client_email: 'test@my-gcp-project.iam.gserviceaccount.com',
308+
client_id: '123456789',
309+
auth_uri: 'https://accounts.google.com/o/oauth2/auth',
310+
token_uri: 'https://oauth2.googleapis.com/token'
311+
});
312+
const bigqueryConfig: DatabaseIntegrationConfig = {
313+
id: 'my-bigquery',
314+
name: 'Analytics BQ',
315+
type: 'big-query',
316+
metadata: {
317+
authMethod: 'service-account',
318+
service_account: serviceAccountJson
319+
}
320+
};
321+
const project = createMockProject('project-123', [
322+
{ id: 'my-bigquery', name: 'Analytics BQ', type: 'big-query' }
323+
]);
324+
325+
when(notebook.metadata).thenReturn({ deepnoteProjectId: 'project-123' });
326+
when(notebookEditorProvider.findAssociatedNotebookDocument(resource)).thenReturn(instance(notebook));
327+
when(notebookManager.getOriginalProject('project-123')).thenReturn(project);
328+
when(integrationStorage.getIntegrationConfig('my-bigquery')).thenResolve(bigqueryConfig);
329+
330+
const result = await provider.getEnvironmentVariables(resource);
331+
332+
const expectedEnvVarName = 'SQL_MY_BIGQUERY';
333+
assert.ok(result[expectedEnvVarName], `Should have ${expectedEnvVarName} env var`);
334+
335+
const envVarValue = result[expectedEnvVarName];
336+
assert.ok(typeof envVarValue === 'string', 'Env var value should be a string');
337+
assert.ok(envVarValue, 'Env var value should not be undefined');
338+
339+
// Parse and verify the structure
340+
const parsed = JSON.parse(envVarValue!);
341+
// BigQuery env vars should contain connection details
342+
// The exact structure depends on the database-integrations library
343+
assert.ok(parsed, 'Should have parsed BigQuery config');
344+
});
345+
346+
test('DuckDB (dataframe-sql) integration is always included', async () => {
347+
const resource = Uri.file('/test/notebook.deepnote');
348+
const notebook = mock<NotebookDocument>();
349+
const project = createMockProject('project-123', []);
350+
351+
when(notebook.metadata).thenReturn({ deepnoteProjectId: 'project-123' });
352+
when(notebookEditorProvider.findAssociatedNotebookDocument(resource)).thenReturn(instance(notebook));
353+
when(notebookManager.getOriginalProject('project-123')).thenReturn(project);
354+
355+
const result = await provider.getEnvironmentVariables(resource);
356+
357+
// DuckDB integration should generate an env var
358+
// The exact name depends on DATAFRAME_SQL_INTEGRATION_ID
359+
const expectedEnvVarName = `SQL_${DATAFRAME_SQL_INTEGRATION_ID.toUpperCase().replace(/-/g, '_')}`;
360+
assert.ok(result[expectedEnvVarName], `Should have ${expectedEnvVarName} env var for DuckDB`);
361+
});
362+
});
363+
});
364+
365+
suite('onDidChangeEnvironmentVariables event', () => {
366+
test('Fires when integration storage changes', (done) => {
367+
let eventFired = false;
368+
provider.onDidChangeEnvironmentVariables(() => {
369+
eventFired = true;
370+
assert.ok(true, 'Event should fire when integrations change');
371+
done();
372+
});
373+
374+
// Trigger the integration storage change event
375+
onDidChangeIntegrationsEmitter.fire();
376+
377+
// Give it a moment to propagate
378+
setTimeout(() => {
379+
if (!eventFired) {
380+
done(new Error('Event did not fire'));
381+
}
382+
}, 100);
383+
});
384+
});
385+
});

0 commit comments

Comments
 (0)