Skip to content

Commit c9a64e4

Browse files
jankucajankuca
committed
test: add tests for snowflake integrations
1 parent d664998 commit c9a64e4

File tree

1 file changed

+300
-1
lines changed

1 file changed

+300
-1
lines changed

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

Lines changed: 300 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,13 @@ import { CancellationTokenSource, EventEmitter, NotebookCell, NotebookCellKind,
55
import { IDisposableRegistry } from '../../common/types';
66
import { IntegrationStorage } from './integrationStorage';
77
import { SqlIntegrationEnvironmentVariablesProvider } from './sqlIntegrationEnvironmentVariablesProvider';
8-
import { IntegrationType, PostgresIntegrationConfig, BigQueryIntegrationConfig } from './integrationTypes';
8+
import {
9+
IntegrationType,
10+
PostgresIntegrationConfig,
11+
BigQueryIntegrationConfig,
12+
SnowflakeIntegrationConfig,
13+
SnowflakeAuthMethods
14+
} from './integrationTypes';
915
import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock';
1016

1117
suite('SqlIntegrationEnvironmentVariablesProvider', () => {
@@ -366,6 +372,299 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => {
366372
const envVars = await provider.getEnvironmentVariables(uri, cts.token);
367373
assert.deepStrictEqual(envVars, {});
368374
});
375+
376+
suite('Snowflake Integration', () => {
377+
test('Returns environment variable for Snowflake with PASSWORD auth', async () => {
378+
const uri = Uri.file('/test/notebook.deepnote');
379+
const integrationId = 'my-snowflake';
380+
const config: SnowflakeIntegrationConfig = {
381+
id: integrationId,
382+
name: 'My Snowflake',
383+
type: IntegrationType.Snowflake,
384+
account: 'myorg-myaccount',
385+
warehouse: 'COMPUTE_WH',
386+
database: 'MYDB',
387+
role: 'ANALYST',
388+
authMethod: SnowflakeAuthMethods.PASSWORD,
389+
username: 'john.doe',
390+
password: 'secret123'
391+
};
392+
393+
const notebook = createMockNotebook(uri, [
394+
createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT * FROM customers', {
395+
sql_integration_id: integrationId
396+
})
397+
]);
398+
399+
when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]);
400+
when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config);
401+
402+
const envVars = await provider.getEnvironmentVariables(uri);
403+
404+
assert.property(envVars, 'SQL_MY_SNOWFLAKE');
405+
const credentialsJson = JSON.parse(envVars['SQL_MY_SNOWFLAKE']!);
406+
assert.strictEqual(
407+
credentialsJson.url,
408+
'snowflake://john.doe:secret123@myorg-myaccount/MYDB?warehouse=COMPUTE_WH&role=ANALYST&application=Deepnote'
409+
);
410+
assert.deepStrictEqual(credentialsJson.params, {});
411+
assert.strictEqual(credentialsJson.param_style, 'format');
412+
});
413+
414+
test('Returns environment variable for Snowflake with legacy null auth (username+password)', async () => {
415+
const uri = Uri.file('/test/notebook.deepnote');
416+
const integrationId = 'legacy-snowflake';
417+
const config: SnowflakeIntegrationConfig = {
418+
id: integrationId,
419+
name: 'Legacy Snowflake',
420+
type: IntegrationType.Snowflake,
421+
account: 'legacy-account',
422+
warehouse: 'WH',
423+
database: 'DB',
424+
authMethod: null,
425+
username: 'user',
426+
password: 'pass'
427+
};
428+
429+
const notebook = createMockNotebook(uri, [
430+
createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT 1', {
431+
sql_integration_id: integrationId
432+
})
433+
]);
434+
435+
when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]);
436+
when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config);
437+
438+
const envVars = await provider.getEnvironmentVariables(uri);
439+
440+
assert.property(envVars, 'SQL_LEGACY_SNOWFLAKE');
441+
const credentialsJson = JSON.parse(envVars['SQL_LEGACY_SNOWFLAKE']!);
442+
assert.strictEqual(
443+
credentialsJson.url,
444+
'snowflake://user:pass@legacy-account/DB?warehouse=WH&application=Deepnote'
445+
);
446+
assert.deepStrictEqual(credentialsJson.params, {});
447+
});
448+
449+
test('Returns environment variable for Snowflake with SERVICE_ACCOUNT_KEY_PAIR auth', async () => {
450+
const uri = Uri.file('/test/notebook.deepnote');
451+
const integrationId = 'snowflake-keypair';
452+
const privateKey = '-----BEGIN PRIVATE KEY-----\nMIIEvQIBADANBg...\n-----END PRIVATE KEY-----';
453+
const config: SnowflakeIntegrationConfig = {
454+
id: integrationId,
455+
name: 'Snowflake KeyPair',
456+
type: IntegrationType.Snowflake,
457+
account: 'keypair-account',
458+
warehouse: 'ETL_WH',
459+
database: 'PROD_DB',
460+
role: 'ETL_ROLE',
461+
authMethod: SnowflakeAuthMethods.SERVICE_ACCOUNT_KEY_PAIR,
462+
username: 'service_account',
463+
privateKey: privateKey,
464+
privateKeyPassphrase: 'passphrase123'
465+
};
466+
467+
const notebook = createMockNotebook(uri, [
468+
createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT * FROM events', {
469+
sql_integration_id: integrationId
470+
})
471+
]);
472+
473+
when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]);
474+
when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config);
475+
476+
const envVars = await provider.getEnvironmentVariables(uri);
477+
478+
assert.property(envVars, 'SQL_SNOWFLAKE_KEYPAIR');
479+
const credentialsJson = JSON.parse(envVars['SQL_SNOWFLAKE_KEYPAIR']!);
480+
assert.strictEqual(
481+
credentialsJson.url,
482+
'snowflake://service_account@keypair-account/PROD_DB?warehouse=ETL_WH&role=ETL_ROLE&authenticator=snowflake_jwt&application=Deepnote'
483+
);
484+
assert.deepStrictEqual(credentialsJson.params, {
485+
private_key: privateKey,
486+
private_key_passphrase: 'passphrase123'
487+
});
488+
assert.strictEqual(credentialsJson.param_style, 'format');
489+
});
490+
491+
test('Returns environment variable for Snowflake with SERVICE_ACCOUNT_KEY_PAIR auth without passphrase', async () => {
492+
const uri = Uri.file('/test/notebook.deepnote');
493+
const integrationId = 'snowflake-keypair-no-pass';
494+
const privateKey = '-----BEGIN PRIVATE KEY-----\nMIIEvQIBADANBg...\n-----END PRIVATE KEY-----';
495+
const config: SnowflakeIntegrationConfig = {
496+
id: integrationId,
497+
name: 'Snowflake KeyPair No Pass',
498+
type: IntegrationType.Snowflake,
499+
account: 'account123',
500+
warehouse: 'WH',
501+
database: 'DB',
502+
authMethod: SnowflakeAuthMethods.SERVICE_ACCOUNT_KEY_PAIR,
503+
username: 'svc_user',
504+
privateKey: privateKey
505+
};
506+
507+
const notebook = createMockNotebook(uri, [
508+
createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT 1', {
509+
sql_integration_id: integrationId
510+
})
511+
]);
512+
513+
when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]);
514+
when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config);
515+
516+
const envVars = await provider.getEnvironmentVariables(uri);
517+
518+
assert.property(envVars, 'SQL_SNOWFLAKE_KEYPAIR_NO_PASS');
519+
const credentialsJson = JSON.parse(envVars['SQL_SNOWFLAKE_KEYPAIR_NO_PASS']!);
520+
assert.strictEqual(
521+
credentialsJson.url,
522+
'snowflake://svc_user@account123/DB?warehouse=WH&authenticator=snowflake_jwt&application=Deepnote'
523+
);
524+
assert.deepStrictEqual(credentialsJson.params, {
525+
private_key: privateKey
526+
});
527+
});
528+
529+
test('Properly encodes special characters in Snowflake credentials', async () => {
530+
const uri = Uri.file('/test/notebook.deepnote');
531+
const integrationId = 'snowflake-special';
532+
const config: SnowflakeIntegrationConfig = {
533+
id: integrationId,
534+
name: 'Snowflake Special',
535+
type: IntegrationType.Snowflake,
536+
account: 'my-org.account',
537+
warehouse: 'WH@2024',
538+
database: 'DB:TEST',
539+
role: 'ROLE#1',
540+
authMethod: SnowflakeAuthMethods.PASSWORD,
541+
username: '[email protected]',
542+
password: 'p@ss:word!#$%'
543+
};
544+
545+
const notebook = createMockNotebook(uri, [
546+
createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT 1', {
547+
sql_integration_id: integrationId
548+
})
549+
]);
550+
551+
when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]);
552+
when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config);
553+
554+
const envVars = await provider.getEnvironmentVariables(uri);
555+
556+
assert.property(envVars, 'SQL_SNOWFLAKE_SPECIAL');
557+
const credentialsJson = JSON.parse(envVars['SQL_SNOWFLAKE_SPECIAL']!);
558+
// Verify URL encoding of special characters
559+
assert.strictEqual(
560+
credentialsJson.url,
561+
'snowflake://user%40domain.com:p%40ss%3Aword!%23%24%[email protected]/DB%3ATEST?warehouse=WH%402024&role=ROLE%231&application=Deepnote'
562+
);
563+
});
564+
565+
test('Handles Snowflake with minimal optional fields', async () => {
566+
const uri = Uri.file('/test/notebook.deepnote');
567+
const integrationId = 'snowflake-minimal';
568+
const config: SnowflakeIntegrationConfig = {
569+
id: integrationId,
570+
name: 'Snowflake Minimal',
571+
type: IntegrationType.Snowflake,
572+
account: 'minimal-account',
573+
authMethod: SnowflakeAuthMethods.PASSWORD,
574+
username: 'user',
575+
password: 'pass'
576+
};
577+
578+
const notebook = createMockNotebook(uri, [
579+
createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT 1', {
580+
sql_integration_id: integrationId
581+
})
582+
]);
583+
584+
when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]);
585+
when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config);
586+
587+
const envVars = await provider.getEnvironmentVariables(uri);
588+
589+
assert.property(envVars, 'SQL_SNOWFLAKE_MINIMAL');
590+
const credentialsJson = JSON.parse(envVars['SQL_SNOWFLAKE_MINIMAL']!);
591+
// Should not include warehouse, database, or role in URL when not provided
592+
assert.strictEqual(credentialsJson.url, 'snowflake://user:pass@minimal-account?application=Deepnote');
593+
});
594+
595+
test('Throws error for unsupported Snowflake auth method (OKTA)', async () => {
596+
const uri = Uri.file('/test/notebook.deepnote');
597+
const integrationId = 'snowflake-okta';
598+
const config: SnowflakeIntegrationConfig = {
599+
id: integrationId,
600+
name: 'Snowflake OKTA',
601+
type: IntegrationType.Snowflake,
602+
account: 'okta-account',
603+
authMethod: SnowflakeAuthMethods.OKTA
604+
};
605+
606+
const notebook = createMockNotebook(uri, [
607+
createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT 1', {
608+
sql_integration_id: integrationId
609+
})
610+
]);
611+
612+
when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]);
613+
when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config);
614+
615+
// Should return empty object when unsupported auth method is encountered
616+
const envVars = await provider.getEnvironmentVariables(uri);
617+
assert.deepStrictEqual(envVars, {});
618+
});
619+
620+
test('Throws error for unsupported Snowflake auth method (AZURE_AD)', async () => {
621+
const uri = Uri.file('/test/notebook.deepnote');
622+
const integrationId = 'snowflake-azure';
623+
const config: SnowflakeIntegrationConfig = {
624+
id: integrationId,
625+
name: 'Snowflake Azure',
626+
type: IntegrationType.Snowflake,
627+
account: 'azure-account',
628+
authMethod: SnowflakeAuthMethods.AZURE_AD
629+
};
630+
631+
const notebook = createMockNotebook(uri, [
632+
createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT 1', {
633+
sql_integration_id: integrationId
634+
})
635+
]);
636+
637+
when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]);
638+
when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config);
639+
640+
const envVars = await provider.getEnvironmentVariables(uri);
641+
assert.deepStrictEqual(envVars, {});
642+
});
643+
644+
test('Throws error for unsupported Snowflake auth method (KEY_PAIR)', async () => {
645+
const uri = Uri.file('/test/notebook.deepnote');
646+
const integrationId = 'snowflake-keypair-user';
647+
const config: SnowflakeIntegrationConfig = {
648+
id: integrationId,
649+
name: 'Snowflake KeyPair User',
650+
type: IntegrationType.Snowflake,
651+
account: 'keypair-user-account',
652+
authMethod: SnowflakeAuthMethods.KEY_PAIR
653+
};
654+
655+
const notebook = createMockNotebook(uri, [
656+
createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT 1', {
657+
sql_integration_id: integrationId
658+
})
659+
]);
660+
661+
when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]);
662+
when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config);
663+
664+
const envVars = await provider.getEnvironmentVariables(uri);
665+
assert.deepStrictEqual(envVars, {});
666+
});
667+
});
369668
});
370669

371670
function createMockNotebook(uri: Uri, cells: NotebookCell[]): NotebookDocument {

0 commit comments

Comments
 (0)