test: cover new logic

Author: jankuca

src/notebooks/deepnote/deepnoteNotebookManager.unit.test.ts

@@ -183,6 +183,83 @@ suite('DeepnoteNotebookManager', () => {
         });
     });
 
+    suite('updateProjectIntegrations', () => {
+        test('should update integrations list for existing project', () => {
+            manager.storeOriginalProject('project-123', mockProject, 'notebook-456');
+
+            const integrations = [
+                { id: 'int-1', name: 'PostgreSQL', type: 'pgsql' },
+                { id: 'int-2', name: 'BigQuery', type: 'big-query' }
+            ];
+
+            manager.updateProjectIntegrations('project-123', integrations);
+
+            const updatedProject = manager.getOriginalProject('project-123');
+            assert.deepStrictEqual(updatedProject?.project.integrations, integrations);
+        });
+
+        test('should replace existing integrations list', () => {
+            const projectWithIntegrations: DeepnoteProject = {
+                ...mockProject,
+                project: {
+                    ...mockProject.project,
+                    integrations: [{ id: 'old-int', name: 'Old Integration', type: 'pgsql' }]
+                }
+            };
+
+            manager.storeOriginalProject('project-123', projectWithIntegrations, 'notebook-456');
+
+            const newIntegrations = [
+                { id: 'new-int-1', name: 'New Integration 1', type: 'pgsql' },
+                { id: 'new-int-2', name: 'New Integration 2', type: 'big-query' }
+            ];
+
+            manager.updateProjectIntegrations('project-123', newIntegrations);
+
+            const updatedProject = manager.getOriginalProject('project-123');
+            assert.deepStrictEqual(updatedProject?.project.integrations, newIntegrations);
+        });
+
+        test('should handle empty integrations array', () => {
+            const projectWithIntegrations: DeepnoteProject = {
+                ...mockProject,
+                project: {
+                    ...mockProject.project,
+                    integrations: [{ id: 'int-1', name: 'Integration 1', type: 'pgsql' }]
+                }
+            };
+
+            manager.storeOriginalProject('project-123', projectWithIntegrations, 'notebook-456');
+
+            manager.updateProjectIntegrations('project-123', []);
+
+            const updatedProject = manager.getOriginalProject('project-123');
+            assert.deepStrictEqual(updatedProject?.project.integrations, []);
+        });
+
+        test('should do nothing for unknown project', () => {
+            // Should not throw an error
+            manager.updateProjectIntegrations('unknown-project', [{ id: 'int-1', name: 'Integration', type: 'pgsql' }]);
+
+            const project = manager.getOriginalProject('unknown-project');
+            assert.strictEqual(project, undefined);
+        });
+
+        test('should preserve other project properties', () => {
+            manager.storeOriginalProject('project-123', mockProject, 'notebook-456');
+
+            const integrations = [{ id: 'int-1', name: 'PostgreSQL', type: 'pgsql' }];
+
+            manager.updateProjectIntegrations('project-123', integrations);
+
+            const updatedProject = manager.getOriginalProject('project-123');
+            assert.strictEqual(updatedProject?.project.id, mockProject.project.id);
+            assert.strictEqual(updatedProject?.project.name, mockProject.project.name);
+            assert.strictEqual(updatedProject?.version, mockProject.version);
+            assert.deepStrictEqual(updatedProject?.metadata, mockProject.metadata);
+        });
+    });
+
     suite('integration scenarios', () => {
         test('should handle complete workflow for multiple projects', () => {
             manager.storeOriginalProject('project-1', mockProject, 'notebook-1');

src/notebooks/deepnote/sqlCellStatusBarProvider.unit.test.ts

@@ -701,6 +701,72 @@ suite('SqlCellStatusBarProvider', () => {
             assert.isDefined(currentItem, 'Current integration should be in quick pick items');
             assert.strictEqual(currentItem.detail, 'Currently selected');
         });
+
+        test('shows error message when project ID is missing', async () => {
+            const cell = createMockCell('sql', {}, {}); // No notebook metadata
+
+            await switchIntegrationHandler(cell);
+
+            verify(mockedVSCodeNamespaces.window.showErrorMessage(anything())).once();
+            verify(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).never();
+        });
+
+        test('shows error message when project is not found', async () => {
+            const notebookMetadata = { deepnoteProjectId: 'missing-project' };
+            const cell = createMockCell('sql', {}, notebookMetadata);
+
+            when(commandNotebookManager.getOriginalProject('missing-project')).thenReturn(undefined);
+
+            await switchIntegrationHandler(cell);
+
+            verify(mockedVSCodeNamespaces.window.showErrorMessage(anything())).once();
+            verify(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).never();
+        });
+
+        test('skips DATAFRAME_SQL_INTEGRATION_ID from project integrations list', async () => {
+            const notebookMetadata = { deepnoteProjectId: 'project-1' };
+            const cell = createMockCell('sql', {}, notebookMetadata);
+            let quickPickItems: any[] = [];
+
+            when(commandNotebookManager.getOriginalProject('project-1')).thenReturn({
+                project: {
+                    integrations: [
+                        {
+                            id: DATAFRAME_SQL_INTEGRATION_ID,
+                            name: 'Should be skipped',
+                            type: 'duckdb'
+                        },
+                        {
+                            id: 'postgres-integration',
+                            name: 'PostgreSQL',
+                            type: 'pgsql'
+                        }
+                    ]
+                }
+            } as any);
+            when(mockedVSCodeNamespaces.window.showErrorMessage(anything())).thenReturn(Promise.resolve(undefined));
+            when(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).thenCall((items) => {
+                quickPickItems = items;
+                return Promise.resolve(undefined);
+            });
+
+            await switchIntegrationHandler(cell);
+
+            // Should have 2 items: postgres-integration and DuckDB (added separately)
+            const projectIntegrationItems = quickPickItems.filter(
+                (item) => item.id && item.id !== DATAFRAME_SQL_INTEGRATION_ID
+            );
+            assert.strictEqual(
+                projectIntegrationItems.length,
+                1,
+                'Should have only 1 project integration (DATAFRAME_SQL_INTEGRATION_ID should be skipped)'
+            );
+            assert.strictEqual(projectIntegrationItems[0].id, 'postgres-integration');
+
+            // DuckDB should still be in the list (added separately)
+            const duckDbItem = quickPickItems.find((item) => item.id === DATAFRAME_SQL_INTEGRATION_ID);
+            assert.isDefined(duckDbItem, 'DuckDB should still be in the list');
+        });
     });
 
     function createMockCell(

src/platform/notebooks/deepnote/integrationTypes.ts

@@ -15,27 +15,25 @@ export enum IntegrationType {
 /**
  * Map Deepnote integration type strings to our IntegrationType enum
  */
+const DEEPNOTE_TO_INTEGRATION_TYPE: Record<string, IntegrationType> = {
+    pgsql: IntegrationType.Postgres,
+    'big-query': IntegrationType.BigQuery
+};
+
 export function mapDeepnoteIntegrationType(deepnoteType: string): IntegrationType | undefined {
-    switch (deepnoteType) {
-        case 'pgsql':
-            return IntegrationType.Postgres;
-        case 'big-query':
-            return IntegrationType.BigQuery;
-        default:
-            return undefined;
-    }
+    return DEEPNOTE_TO_INTEGRATION_TYPE[deepnoteType];
 }
 
 /**
  * Map our IntegrationType enum to Deepnote integration type strings
  */
+const INTEGRATION_TYPE_TO_DEEPNOTE: Record<IntegrationType, string> = {
+    [IntegrationType.Postgres]: 'pgsql',
+    [IntegrationType.BigQuery]: 'big-query'
+};
+
 export function mapToDeepnoteIntegrationType(type: IntegrationType): string {
-    switch (type) {
-        case IntegrationType.Postgres:
-            return 'pgsql';
-        case IntegrationType.BigQuery:
-            return 'big-query';
-    }
+    return INTEGRATION_TYPE_TO_DEEPNOTE[type];
 }
 
 /**