From 9ee8c029573a3ce47bd5f103da9feaf8b6f62c53 Mon Sep 17 00:00:00 2001
From: Alberto Schiabel <jkomyno@users.noreply.github.com>
Date: Fri, 7 Jun 2024 14:41:18 +0200
Subject: [PATCH] test(query-engine): fix some PlanetScale tests (#4909)

* test(query-engine): Fix PlanetScale test, #4482

* test(query-engine): Fix PlanetScale test, #4485
---
 .../compound_fks_mixed_requiredness.rs        | 25 +++++++++++--------
 .../writes/top_level_mutations/update_many.rs | 12 ++++-----
 2 files changed, 21 insertions(+), 16 deletions(-)

diff --git a/query-engine/connector-test-kit-rs/query-engine-tests/tests/writes/relations/compound_fks_mixed_requiredness.rs b/query-engine/connector-test-kit-rs/query-engine-tests/tests/writes/relations/compound_fks_mixed_requiredness.rs
index d1be84d2d9a9..c0c80a5a8e14 100644
--- a/query-engine/connector-test-kit-rs/query-engine-tests/tests/writes/relations/compound_fks_mixed_requiredness.rs
+++ b/query-engine/connector-test-kit-rs/query-engine-tests/tests/writes/relations/compound_fks_mixed_requiredness.rs
@@ -25,11 +25,11 @@ mod compound_fks {
         schema.to_owned()
     }
 
-    // "A One to Many relation with mixed requiredness" should "be writable and readable"-
-    // In PlanetScale, this fails with:
-    // `Expected result to return an error, but found success: {"data":{"createOnePost":{"id":2,"user_id":2,"user_age":2,"User":null}}}`
-    #[connector_test(exclude(MySql(5.6), MongoDb, Vitess("planetscale.js", "planetscale.js.wasm")))]
+    // "A One to Many relation with mixed requiredness" should "be writable and readable"
+    #[connector_test(exclude(MySql(5.6), MongoDb))]
     async fn one2m_mix_required_writable_readable(runner: Runner) -> TestResult<()> {
+        use query_tests_setup::{ConnectorVersion, VitessVersion::*};
+
         // Setup user
         insta::assert_snapshot!(
           run_query!(&runner, r#"mutation{createOneUser(data:{id: 1, nr:1, age: 1}){id, nr, age, Post{id}}}"#),
@@ -49,12 +49,17 @@ mod compound_fks {
           @r###"{"data":{"createOnePost":{"id":1,"user_id":1,"user_age":null,"User":null}}}"###
         );
 
-        // Foreign key violation
-        assert_error!(
-            &runner,
-            r#"mutation{createOnePost(data:{id: 2, user_id:2, user_age: 2}){id, user_id, user_age, User{id}}}"#,
-            2003
-        );
+        // Foreign key violation, which doesn't happen on PlanetScale.
+        if !matches!(
+            runner.connector_version(),
+            ConnectorVersion::Vitess(Some(PlanetscaleJsNapi)) | ConnectorVersion::Vitess(Some(PlanetscaleJsWasm))
+        ) {
+            assert_error!(
+                &runner,
+                r#"mutation{createOnePost(data:{id: 2, user_id:2, user_age: 2}){id, user_id, user_age, User{id}}}"#,
+                2003
+            );
+        }
 
         // Success
         insta::assert_snapshot!(
diff --git a/query-engine/connector-test-kit-rs/query-engine-tests/tests/writes/top_level_mutations/update_many.rs b/query-engine/connector-test-kit-rs/query-engine-tests/tests/writes/top_level_mutations/update_many.rs
index f821bb43dcab..c9a6d6df567c 100644
--- a/query-engine/connector-test-kit-rs/query-engine-tests/tests/writes/top_level_mutations/update_many.rs
+++ b/query-engine/connector-test-kit-rs/query-engine-tests/tests/writes/top_level_mutations/update_many.rs
@@ -123,7 +123,7 @@ mod update_many {
     }
 
     // "An updateMany mutation" should "correctly apply all number operations for Int"
-    #[connector_test(exclude(Vitess("planetscale.js", "planetscale.js.wasm"), CockroachDb))]
+    #[connector_test(exclude(CockroachDb))]
     async fn apply_number_ops_for_int(runner: Runner) -> TestResult<()> {
         create_row(&runner, r#"{ id: 1, optStr: "str1" }"#).await?;
         create_row(&runner, r#"{ id: 2, optStr: "str2", optInt: 2 }"#).await?;
@@ -240,7 +240,7 @@ mod update_many {
     }
 
     // "An updateMany mutation" should "correctly apply all number operations for Float"
-    #[connector_test(exclude(Vitess("planetscale.js", "planetscale.js.wasm")))]
+    #[connector_test]
     async fn apply_number_ops_for_float(runner: Runner) -> TestResult<()> {
         create_row(&runner, r#"{ id: 1, optStr: "str1" }"#).await?;
         create_row(&runner, r#"{ id: 2, optStr: "str2", optFloat: 2 }"#).await?;
@@ -297,10 +297,10 @@ mod update_many {
         let count = &res["data"]["updateManyTestModel"]["count"];
 
         // MySql does not count incrementing a null so the count is different
-        if !matches!(runner.connector_version(), ConnectorVersion::MySql(_)) {
-            // On PlanetScale, this fails with:
-            //   left: Number(2)
-            //   right: 3
+        if !matches!(
+            runner.connector_version(),
+            ConnectorVersion::MySql(_) | ConnectorVersion::Vitess(_)
+        ) {
             assert_eq!(count, 3);
         }