diff --git a/nixos/lib/test-driver/Machine.pm b/nixos/lib/test-driver/Machine.pm
index a7ed5d1faa38..dbed242eb865 100644
--- a/nixos/lib/test-driver/Machine.pm
+++ b/nixos/lib/test-driver/Machine.pm
@@ -602,8 +602,11 @@ sub waitForX {
     $self->nest("waiting for the X11 server", sub {
         retry sub {
             my ($status, $out) = $self->execute("journalctl -b SYSLOG_IDENTIFIER=systemd | grep 'session opened'");
+            my $display = $ENV{'DISPLAY'};
+            $display =~ s/://;
+
             return 0 if $status != 0;
-            ($status, $out) = $self->execute("[ -e /tmp/.X11-unix/X0 ]");
+            ($status, $out) = $self->execute("[ -e /tmp/.X11-unix/X$display ]");
             return 1 if $status == 0;
         }
     });
diff --git a/nixos/modules/testing/test-instrumentation.nix b/nixos/modules/testing/test-instrumentation.nix
index 7f5b55d5cca0..2465e99e9c16 100644
--- a/nixos/modules/testing/test-instrumentation.nix
+++ b/nixos/modules/testing/test-instrumentation.nix
@@ -5,9 +5,10 @@
 
 with lib;
 
-let kernel = config.boot.kernelPackages.kernel; in
-
-{
+let
+  display = (d: if d == null then 0 else d) config.services.xserver.display;
+  kernel = config.boot.kernelPackages.kernel;
+in {
 
   # This option is a dummy that if used in conjunction with
   # modules/virtualisation/qemu-vm.nix gets merged with the same option defined
@@ -28,7 +29,7 @@ let kernel = config.boot.kernelPackages.kernel; in
           ''
             export USER=root
             export HOME=/root
-            export DISPLAY=:0.0
+            export DISPLAY=:${toString display}
 
             source /etc/profile