test(cli): support datasource extras in test utils

Enhances the test utility helpers to allow passing extra datasource properties, such as multi-schema configurations for PostgreSQL.

Refactors existing database pull tests to use these extra properties, ensuring the generated ZModel schema correctly reflects multi-schema environments while simplifying assertions.

Author: svetch

packages/cli/src/actions/pull/utils.ts

@@ -64,12 +64,11 @@ export function getDatasource(model: Model) {
 
     const schemasField = datasource.fields.find((f) => f.name === 'schemas');
     const schemas =
-        (schemasField &&
-            getLiteralArray(schemasField.value)
-                ?.map(getStringLiteral)
-                .filter((s) => s !== undefined)) ||
+    (schemasField &&
+        getLiteralArray(schemasField.value)
+        ?.filter((s) => s !== undefined)) as string[] ||
         [];
-
+    
     return {
         name: datasource.name,
         provider: getStringLiteral(

packages/cli/test/db/pull.test.ts

@@ -353,21 +353,19 @@ model Post {
 
     @@schema("content")
 }`,
-            { provider: 'postgresql' },
+            { provider: 'postgresql', extra:{ schemas: ["public", "content", "auth"] } },
         );
         runCli('db push', workDir);
 
         const schemaFile = path.join(workDir, 'zenstack/schema.zmodel');
         const { model } = await loadSchemaDocument(schemaFile, { returnServices: true });
         const expectedSchema = generator.generate(model);
 
-        fs.writeFileSync(schemaFile, getDefaultPrelude({ provider: 'postgresql' }));
+        fs.writeFileSync(schemaFile, getDefaultPrelude({ provider: 'postgresql', extra:{ schemas: ["public", "content", "auth"]} }));
         runCli('db pull --indent 4', workDir);
 
         const restoredSchema = getSchema(workDir);
         expect(restoredSchema).toEqual(expectedSchema);
-        expect(restoredSchema).toContain('@@schema("auth")');
-        expect(restoredSchema).toContain('@@schema("content")');
     });
 
     it('should preserve native PostgreSQL enums when schema exists', async ({ skip }) => {
@@ -404,8 +402,6 @@ enum UserRole {
         const pulledSchema = getSchema(workDir);
 
         expect(pulledSchema).toEqual(originalSchema);
-        expect(pulledSchema).toContain('enum UserStatus');
-        expect(pulledSchema).toContain('enum UserRole');
     });
 
     it('should not modify schema with PostgreSQL-specific features', async ({ skip }) => {
@@ -442,7 +438,7 @@ enum UserStatus {
     INACTIVE
     SUSPENDED
 }`,
-            { provider: 'postgresql' },
+            { provider: 'postgresql', extra:{ schemas: ["public", "content", "auth"] }  },
         );
         runCli('db push', workDir);
 

packages/cli/test/utils.ts

@@ -42,7 +42,7 @@ function getTestDbName(provider: string) {
     );
 }
 
-export function getDefaultPrelude(options?: { provider?: 'sqlite' | 'postgresql' | 'mysql' }) {
+export function getDefaultPrelude(options?: { provider?: 'sqlite' | 'postgresql' | 'mysql', extra?: Record<string, string | string[]> }) {
     const provider = (options?.provider || getTestDbProvider()) ?? 'sqlite';
     const dbName = getTestDbName(provider);
     let dbUrl: string;
@@ -60,10 +60,10 @@ export function getDefaultPrelude(options?: { provider?: 'sqlite' | 'postgresql'
         default:
             throw new Error(`Unsupported provider: ${provider}`);
     }
-
+    const DATASOURCE_EXTRAS = Object.entries(options?.extra || {}).map(([k, v]) => `  ${k} = ${JSON.stringify(v)}`).join('\n');
     const ZMODEL_PRELUDE = `datasource db {
   provider = "${provider}"
-  url      = "${dbUrl}"
+  url      = "${dbUrl}"${DATASOURCE_EXTRAS.length > 0 ? '\n' + DATASOURCE_EXTRAS : ''}
 }`;
     return ZMODEL_PRELUDE;
 }
@@ -81,9 +81,9 @@ export function createProject(
 
 export async function createFormattedProject(
     zmodel: string,
-    options?: { provider?: 'sqlite' | 'postgresql' | 'mysql' },
+    options?: { provider?: 'sqlite' | 'postgresql' | 'mysql', extra?: Record<string, string | string[]> },
 ) {
-    const fullContent = `${getDefaultPrelude({ provider: options?.provider })}\n\n${zmodel}`;
+    const fullContent = `${getDefaultPrelude({ provider: options?.provider, extra: options?.extra })}\n\n${zmodel}`;
     const formatted = await formatDocument(fullContent);
     return createProject(formatted, { customPrelude: true, provider: options?.provider });
 }