depot/third_party/nixpkgs/nixos/tests/vscode-remote-ssh.nix

import ./make-test-python.nix (
  { lib, ... }@args:
  let
    pkgs = args.pkgs.extend (
      self: super: {
        stdenv = super.stdenv.override {
          config = super.config // {
            allowUnfreePredicate =
              pkg:
              builtins.elem (lib.getName pkg) [
                "vscode"
                "vscode-with-extensions"
                "vscode-extension-ms-vscode-remote-remote-ssh"
              ];
          };
        };
      }
    );

    inherit (import ./ssh-keys.nix pkgs) snakeOilPrivateKey snakeOilPublicKey;

    inherit (pkgs.vscode.passthru) rev vscodeServer;
  in
  {
    name = "vscode-remote-ssh";

    meta = {
      maintainers = [ ];
      timeout = 600;
    };

    nodes =
      let
        serverAddress = "192.168.0.2";
        clientAddress = "192.168.0.1";
      in
      {
        server =
          { ... }:
          {
            networking.interfaces.eth1.ipv4.addresses = [
              {
                address = serverAddress;
                prefixLength = 24;
              }
            ];
            services.openssh.enable = true;
            users.users.root.openssh.authorizedKeys.keys = [ snakeOilPublicKey ];
            virtualisation.additionalPaths = with pkgs; [
              patchelf
              bintools
              (lib.getLib stdenv.cc.cc)
            ];
          };
        client =
          { ... }:
          {
            imports = [
              ./common/x11.nix
              ./common/user-account.nix
            ];
            networking.interfaces.eth1.ipv4.addresses = [
              {
                address = clientAddress;
                prefixLength = 24;
              }
            ];
            networking.hosts.${serverAddress} = [ "server" ];
            test-support.displayManager.auto.user = "alice";
            environment.systemPackages = [
              (pkgs.vscode-with-extensions.override {
                vscodeExtensions = [
                  pkgs.vscode-extensions.ms-vscode-remote.remote-ssh
                ];
              })
            ];
          };
      };

    enableOCR = true;

    testScript =
      let
        jq = "${pkgs.jq}/bin/jq";

        sshConfig = builtins.toFile "ssh.conf" ''
          UserKnownHostsFile=/dev/null
          StrictHostKeyChecking=no
        '';

        vscodeConfig = builtins.toFile "settings.json" ''
          {
            "window.zoomLevel": 1,
            "security.workspace.trust.startupPrompt": "always"
          }
        '';
      in
      ''
        def connect_with_remote_ssh(screenshot, should_succeed):
          print(f"connect_with_remote_ssh({screenshot=}, {should_succeed=})")

          if server.execute("test -d ~/.vscode-server")[0] == 0:
            server.succeed("rm -r ~/.vscode-server")

          server.succeed("mkdir -p ~/.vscode-server/bin")
          server.succeed("cp -r ${vscodeServer} ~/.vscode-server/bin/${rev}")

          client.succeed("sudo -u alice code --remote=ssh-remote+root@server /root")
          client.wait_for_window("Visual Studio Code")

          if should_succeed:
            ocr_text = "Do you trust"
          else:
            ocr_text = "Could not establish connection"
          client.wait_for_text(ocr_text)
          client.screenshot(screenshot)

          if should_succeed:
            # Press the Don't Trust button
            client.send_key("tab")
            client.send_key("tab")
            client.send_key("tab")
            client.send_key("\n")
          else:
            # Close the error dialog
            client.send_key("esc")

          # Don't send Ctrl-q too quickly otherwise it might not get sent to VS Code
          client.sleep(1)
          client.send_key("ctrl-q")
          client.wait_until_fails("pidof code")


        start_all()
        server.wait_for_open_port(22)

        VSCODE_COMMIT = server.execute("${jq} -r .commit ${pkgs.vscode}/lib/vscode/resources/app/product.json")[1].rstrip()
        SERVER_COMMIT = server.execute("${jq} -r .commit ${vscodeServer}/product.json")[1].rstrip()

        print(f"{VSCODE_COMMIT=} {SERVER_COMMIT=}")
        assert VSCODE_COMMIT == SERVER_COMMIT, "VSCODE_COMMIT and SERVER_COMMIT do not match"

        client.wait_until_succeeds("ping -c1 server")
        client.succeed("sudo -u alice mkdir ~alice/.ssh")
        client.succeed("sudo -u alice install -Dm 600 ${snakeOilPrivateKey} ~alice/.ssh/id_ecdsa")
        client.succeed("sudo -u alice install ${sshConfig} ~alice/.ssh/config")
        client.succeed("sudo -u alice install -Dm 644 ${vscodeConfig} ~alice/.config/Code/User/settings.json")

        client.wait_for_x()
        client.wait_for_file("~alice/.Xauthority")
        client.succeed("xauth merge ~alice/.Xauthority")
        # Move the mouse out of the way
        client.succeed("${pkgs.xdotool}/bin/xdotool mousemove 0 0")

        with subtest("fails to connect when nixpkgs isn't available"):
          server.fail("nix-build '<nixpkgs>' -A hello")
          connect_with_remote_ssh(screenshot="no_node_installed", should_succeed=False)
          server.succeed("test -e ~/.vscode-server/bin/${rev}/node")
          server.fail("~/.vscode-server/bin/${rev}/node -v")

        with subtest("connects when server can patch Node"):
          server.succeed("mkdir -p /nix/var/nix/profiles/per-user/root/channels")
          server.succeed("ln -s ${pkgs.path} /nix/var/nix/profiles/per-user/root/channels/nixos")
          connect_with_remote_ssh(screenshot="build_node_with_nix", should_succeed=True)
      '';
  }
)