query-tests-setup: teardown for external executors

prisma · Sep 5, 2023 · a2729dc · a2729dc
1 parent 095ecf3
commit a2729dc
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 4 deletions.
diff --git a/query-engine/connector-test-kit-rs/query-tests-setup/src/lib.rs b/query-engine/connector-test-kit-rs/query-tests-setup/src/lib.rs
@@ -45,7 +45,12 @@ pub static ENGINE_PROTOCOL: Lazy<String> =
 static NODE_TEST_EXECUTOR: Lazy<Option<String>> = Lazy::new(|| std::env::var("NODE_TEST_EXECUTOR").ok());
 
 /// Teardown of a test setup.
-async fn teardown_project(datamodel: &str, db_schemas: &[&str]) -> TestResult<()> {
+async fn teardown_project(datamodel: &str, db_schemas: &[&str], schema_id: Option<usize>) -> TestResult<()> {
+    if let Some(schema_id) = schema_id {
+        let params = serde_json::json!({ "schemaId": schema_id });
+        crate::executor_process_request::<serde_json::Value>("teardown", params).await?;
+    }
+
     Ok(qe_setup::teardown(datamodel, db_schemas).await?)
 }
 
@@ -169,7 +174,9 @@ fn run_relation_link_test_impl(
 
                 test_fn(&runner, &dm).await.unwrap();
 
-                teardown_project(&datamodel, Default::default()).await.unwrap();
+                teardown_project(&datamodel, Default::default(), runner.schema_id())
+                    .await
+                    .unwrap();
             }
             .with_subscriber(test_tracing_subscriber(
                 ENV_LOG_LEVEL.to_string(),
@@ -277,10 +284,13 @@ fn run_connector_test_impl(
             )
             .await
             .unwrap();
+            let schema_id = runner.schema_id();
 
             test_fn(runner).await.unwrap();
 
-            crate::teardown_project(&datamodel, db_schemas).await.unwrap();
+            crate::teardown_project(&datamodel, db_schemas, schema_id)
+                .await
+                .unwrap();
         }
         .with_subscriber(test_tracing_subscriber(
             ENV_LOG_LEVEL.to_string(),

diff --git a/query-engine/connector-test-kit-rs/query-tests-setup/src/runner/mod.rs b/query-engine/connector-test-kit-rs/query-tests-setup/src/runner/mod.rs
@@ -60,6 +60,13 @@ pub struct Runner {
 }
 
 impl Runner {
+    pub(crate) fn schema_id(&self) -> Option<usize> {
+        match self.executor {
+            RunnerExecutor::Builtin(_) => None,
+            RunnerExecutor::External(schema_id) => Some(schema_id),
+        }
+    }
+
     pub fn prisma_dml(&self) -> &str {
         self.query_schema.internal_data_model.schema.db.source()
     }

diff --git a/query-engine/js-connectors/js/connector-test-kit-executor/src/index.ts b/query-engine/js-connectors/js/connector-test-kit-executor/src/index.ts
@@ -37,7 +37,7 @@ async function handleRequest(method: string, params: unknown): Promise<unknown>
         case 'initializeSchema': {
             interface InitializeSchemaParams {
                 schema: string
-                schemaId: number
+                schemaId: string
                 url: string
             }
 
@@ -63,6 +63,17 @@ async function handleRequest(method: string, params: unknown): Promise<unknown>
 
             return JSON.parse(result)
         }
+        case 'teardown': {
+            interface TeardownPayload {
+                schemaId: number
+            }
+
+            const castParams = params as TeardownPayload;
+            await schemas[castParams.schemaId].disconnect("")
+            delete schemas[castParams.schemaId]
+            return {}
+
+        }
         default: {
             throw new Error(`Unknown method: \`${method}\``)
         }