postgresqlPackages.citus: move test from VM to postgresqlTestExtension

nixos/tests/postgresql/citus.nix

@@ -1,73 +0,0 @@
-{
-  pkgs,
-  makeTest,
-  genTests,
-}:
-
-let
-  inherit (pkgs) lib;
-
-  test-sql = pkgs.writeText "postgresql-test" ''
-    CREATE EXTENSION citus;
-
-    CREATE TABLE examples (
-      id bigserial,
-      shard_key int,
-      PRIMARY KEY (id, shard_key)
-    );
-
-    SELECT create_distributed_table('examples', 'shard_key');
-
-    INSERT INTO examples (shard_key) SELECT shard % 10 FROM generate_series(1,1000) shard;
-  '';
-
-  makeTestFor =
-    package:
-    makeTest {
-      name = "citus-${package.name}";
-      meta = with lib.maintainers; {
-        maintainers = [ typetetris ];
-      };
-
-      nodes.machine =
-        { ... }:
-        {
-          services.postgresql = {
-            inherit package;
-            enable = true;
-            enableJIT = lib.hasInfix "-jit-" package.name;
-            extensions =
-              ps: with ps; [
-                citus
-              ];
-            settings = {
-              shared_preload_libraries = "citus";
-            };
-          };
-        };
-
-      testScript = ''
-        def check_count(statement, lines):
-            return 'test $(sudo -u postgres psql postgres -tAc "{}") -eq {}'.format(
-                statement, lines
-            )
-
-
-        machine.start()
-        machine.wait_for_unit("postgresql")
-
-        with subtest("Postgresql with extension citus is available just after unit start"):
-            machine.succeed(
-                "sudo -u postgres psql -f ${test-sql}"
-            )
-
-        machine.succeed(check_count("SELECT count(*) FROM examples;", 1000))
-
-        machine.shutdown()
-      '';
-    };
-in
-genTests {
-  inherit makeTestFor;
-  filter = _: p: !p.pkgs.citus.meta.broken;
-}

nixos/tests/postgresql/default.nix

@@ -36,7 +36,6 @@ in
 
   # extensions
   anonymizer = importWithArgs ./anonymizer.nix;
-  citus = importWithArgs ./citus.nix;
   pgjwt = importWithArgs ./pgjwt.nix;
   pgvecto-rs = importWithArgs ./pgvecto-rs.nix;
   timescaledb = importWithArgs ./timescaledb.nix;

pkgs/servers/sql/postgresql/ext/citus.nix

@@ -6,6 +6,7 @@
   lz4,
   postgresql,
   postgresqlBuildExtension,
+  postgresqlTestExtension,
   stdenv,
 }:
 
@@ -35,6 +36,33 @@ postgresqlBuildExtension (finalAttrs: {
     lz4
   ];
 
+  passthru.tests.extension = postgresqlTestExtension {
+    inherit (finalAttrs) finalPackage;
+    postgresqlExtraSettings = ''
+      shared_preload_libraries=citus
+    '';
+    sql = ''
+      CREATE EXTENSION citus;
+
+      CREATE TABLE examples (
+        id bigserial,
+        shard_key int,
+        PRIMARY KEY (id, shard_key)
+      );
+
+      SELECT create_distributed_table('examples', 'shard_key');
+
+      INSERT INTO examples (shard_key) SELECT shard % 10 FROM generate_series(1,1000) shard;
+    '';
+    asserts = [
+      {
+        query = "SELECT count(*) FROM examples";
+        expected = "1000";
+        description = "Distributed table can be queried successfully.";
+      }
+    ];
+  };
+
   meta = {
     # "Our soft policy for Postgres version compatibility is to support Citus'
     # latest release with Postgres' 3 latest releases."