ubicloud/bin/regen-screenshots

#!/usr/bin/env ruby
# frozen_string_literal: true

DOCUMENTATION_DIR = "../documentation/"

unless File.directory?(DOCUMENTATION_DIR)
  warn "Documentation site must be checked out in ../documentation"
  exit(1)
end

require "find"
require "capybara"
require "capybara/dsl"
require "capybara/cuprite"
require "puma/cli"
require "nio"
require "securerandom"

ENV["RACK_ENV"] = "test"

# Enable shared connections for Sequel DB, to allow server and screenshot code to share the same transaction
ENV["SHARED_CONNECTION"] = "1"

# Set fake stripe secret key in order to take screenshots of billing pages
ENV["STRIPE_SECRET_KEY"] = "1"

# Set fake GitHub app name in order to take screenshots of GitHub Runner pages
ENV["GITHUB_APP_NAME"] = "1"

require_relative "../loader"

PORT = 8383
db_name = DB.get { current_database.function }
raise "Doesn't look like a test database (#{db_name}), not generating screenshots" unless db_name.end_with?("test")

Capybara.exact = true
Capybara.default_selector = :css
Capybara.default_driver = :cuprite
Capybara.server_port = PORT
Capybara.register_driver(:cuprite) do |app|
  Capybara::Cuprite::Driver.new(app, window_size: [1200, 800], browser_options: {timeout: 15}, base_url: "http://localhost:#{PORT}")
end

require "tilt/erubi"
require "tilt/string"

queue = Queue.new
server = Puma::CLI.new(["-s", "-e", "test", "-b", "tcp://localhost:#{PORT}", "-t", "1:1", "config.ru"])
server.launcher.events.after_booted { queue.push(nil) }
Thread.new do
  server.launcher.run
end
queue.pop

::Mail.defaults do
  delivery_method :test
end

class RegenScreenshots
  include Capybara::DSL

  SCREENSHOTS = {}
  Find.find(DOCUMENTATION_DIR) do |file|
    if File.file?(file) && file.end_with?("screenshot.png")
      SCREENSHOTS[file.delete_prefix(DOCUMENTATION_DIR)] = file
    end
  end

  def screenshot(filename)
    unless (path = SCREENSHOTS.delete(filename))
      raise "No existing screenshot for #{filename} in documentation site"
    end

    # rubocop:disable Lint/Debugger
    save_screenshot(path:)
    # rubocop:enable Lint/Debugger

    puts "Saved screenshot: #{filename}"
  end

  def resize(height, width: 1200)
    Capybara.current_session.driver.browser.resize(width:, height:)
  end

  def call
    visit "/"
    click_link "Create a new account"

    password = SecureRandom.base64(48)
    fill_in "Full Name", with: "Demo"
    fill_in "Email Address", with: "demo@example.com"
    fill_in "Password", with: password
    fill_in "Password Confirmation", with: password
    click_button "Create Account"

    mail = Mail::TestMailer.deliveries.shift
    body = mail.parts[1].decoded
    unless (match = %r{(/verify-account\?key=[^"]+)"}.match(body))
      raise "no verify link in email"
    end
    visit match[1]

    click_button "Verify Account"

    project = Project.first

    resize(900, width: 1600)
    visit "/"
    screenshot "quick-start/managed-services-1-screenshot.png"

    find("#billing-icon").hover
    screenshot "github-actions-integration/quickstart-1-screenshot.png"

    Project.define_method(:has_valid_payment_method?) { true }
    click_link "GitHub Runners"
    screenshot "github-actions-integration/quickstart-2-screenshot.png"

    GithubInstallation.create(installation_id: 123, name: "ubicloud-demo", type: "Organization", project_id: project.id, allocator_preferences: {"family_filter" => ["premium", "standard"]})
    page.refresh
    click_link "Settings", class: "text-gray-500"
    screenshot "github-actions-integration/use-premium-runners-1-screenshot.png"

    resize(650)
    click_link "Tokens"
    screenshot "quick-start/cli-1-screenshot.png"

    click_button "Create Token"
    ApiKey.dataset.update(id: "bf444ee6-2532-8153-975e-af787dbc796e")
    page.refresh
    screenshot "quick-start/cli-2-screenshot.png"

    click_link "PostgreSQL"
    screenshot "managed-postgresql/overview-1-screenshot.png"

    resize(1400)
    click_link "Create PostgreSQL Database"
    screenshot "managed-postgresql/create-screenshot.png"
    screenshot "quick-start/using-kamal-with-ubicloud-3-screenshot.png"

    pg_resource = PostgresResource.create(
      name: "postgresql-demo",
      location_id: Location::HETZNER_FSN1_ID,
      target_vm_size: "standard-2",
      target_storage_size_gib: 10,
      superuser_password: "1",
      project_id: project.id
    ) do |pg|
      pg.id = UBID.parse("pgmjy3v4ef1y7gdpzv6b3fchef").to_uuid
    end
    Project.define_method(:postgres_resources_dataset) do |*a|
      super(*a).with_extend do
        define_method(:all) { [pg_resource] }
        define_method(:first) { |*| pg_resource }
      end
    end

    pg_vm = Vm.create(
      name: "postgresql-demo-vm",
      memory_gib: 8,
      vcpus: 2,
      cores: 2,
      location_id: Location::HETZNER_FSN1_ID,
      project_id: project.id,
      unix_user: "postgres",
      public_key: "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQC7b8ZxEjGHV54sF4z/7H5z9hGJtYI5RvV2kz8KjQhYJvQG8W1vZK8dG9vLWJyZWFrZXItdGVzdEBleGFtcGxlLmNvbYIxP634e6Y0p2FQI+CvLVgNxYqaQZxjx84+SN/XY4FOR4",
      boot_image: "ubuntu-jammy",
      family: "standard"
    )
    pg_timeline = PostgresTimeline.create(
      location_id: Location::HETZNER_FSN1_ID
    ) do |timeline|
      timeline.id = UBID.parse("ptmjy3v4ef1y7gdpzv6b3fchef").to_uuid
    end
    pg_server = PostgresServer.create(
      resource_id: pg_resource.id,
      vm_id: pg_vm.id,
      timeline_id: pg_timeline.id
    ) do |server|
      server.id = UBID.parse("pvmjy3v4ef1y7gdpzv6b3fchef").to_uuid
    end
    PostgresResource.define_method(:display_state) { "running" }
    # Add default firewall rule to the postgres instance
    PostgresFirewallRule.create(postgres_resource_id: pg_resource.id, cidr: "0.0.0.0/0")

    PostgresResource.define_method(:connection_string) { "sample-connection-string" }
    PostgresResource.define_method(:ca_certificates) { "certs" }
    PostgresResource.define_method(:representative_server) { pg_server }
    PostgresResource.define_method(:user_config) { {"max_connections" => 1000, "work_mem" => "64MB"} }
    PostgresResource.define_method(:pgbouncer_user_config) { {"default_pool_size" => 50} }
    Authorization.define_singleton_method(:authorize) { |*| }
    PostgresServer.define_method(:storage_size_gib) { 128 }

    # Create VictoriaMetrics resources for metrics display
    vmr = VictoriaMetricsResource.create(
      name: "victoria-metrics-demo",
      admin_user: "admin",
      admin_password: "password",
      target_vm_size: "standard-2",
      target_storage_size_gib: 100,
      project_id: project.id,
      location_id: Location::HETZNER_FSN1_ID
    )

    vm_metrics = Vm.create(
      name: "victoria-metrics-vm",
      memory_gib: 8,
      vcpus: 2,
      cores: 2,
      location_id: Location::HETZNER_FSN1_ID,
      project_id: project.id,
      unix_user: "ubi",
      public_key: "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQC7b8ZxEjGHV54sF4z/7H5z9hGJtYI5RvV2kz8KjQhYJvQG8W1vZK8dG9vLWJyZWFrZXItdGVzdEBleGFtcGxlLmNvbYIxP634e6Y0p2FQI+CvLVgNxYqaQZxjx84+SN/XY4FOR4",
      boot_image: "ubuntu-jammy",
      family: "standard"
    )

    VictoriaMetricsServer.create(
      victoria_metrics_resource_id: vmr.id,
      vm_id: vm_metrics.id,
      cert: "cert-data",
      cert_key: "cert-key-data"
    )

    # Mock the VictoriaMetrics client to return sample data
    mock_client = Object.new
    def mock_client.query_range(query:, start_ts:, end_ts:)
      # Use a separate RNG for VictoriaMetrics metrics to ensure consistent results
      metrics_rng = Random.new(1234)

      step = 60 # 1 minute intervals
      timestamps = (start_ts..end_ts).step(step).to_a

      case query
      when /node_filesystem_avail_bytes/
        # Return disk usage around 45%
        values = timestamps.map { |ts| [ts, metrics_rng.rand(40..49).to_s] }
        [{
          "labels" => {},
          "values" => values
        }]
      when /node_cpu_seconds_total/
        [{
          "labels" => {"mode" => "user"},
          "values" => timestamps.map { |ts| [ts, metrics_rng.rand(15..64).to_s] }
        },
          {
            "labels" => {"mode" => "system"},
            "values" => timestamps.map { |ts| [ts, metrics_rng.rand(8..12).to_s] }
          },
          {
            "labels" => {"mode" => "iowait"},
            "values" => timestamps.map { |ts| [ts, metrics_rng.rand(2..4).to_s] }
          }]
      else
        # Default sample data
        values = timestamps.map { |ts| [ts, metrics_rng.rand(20..39).to_s] }
        [{
          "labels" => {},
          "values" => values
        }]
      end
    end

    VictoriaMetricsServer.define_method(:client) { |*| mock_client }

    # Mock the metrics_config method on PostgresServer to return the correct project_id
    PostgresServer.define_method(:metrics_config) { {project_id: project.id} }
    Config.define_singleton_method(:postgres_service_project_id) { project.id }

    resize(900, width: 1600)
    click_link "PostgreSQL"
    click_link "postgresql-demo"
    sleep 1 # Wait for metrics charts to load.
    screenshot "managed-postgresql/overview-2-screenshot.png"

    resize(800, width: 1600)
    click_link "Connection"
    screenshot "managed-postgresql/connection-screenshot.png"

    resize(800, width: 1600)
    within("aside nav") do
      click_link "Networking"
    end
    screenshot "managed-postgresql/networking-screenshot.png"

    resize(1200, width: 1600)
    click_link "Resize"
    screenshot "managed-postgresql/resizing-screenshot.png"

    resize(850, width: 1600)
    click_link "Configuration"
    screenshot "managed-postgresql/configuration-screenshot.png"

    resize(1000, width: 1600)
    within("aside nav") do
      click_link "Settings"
    end
    screenshot "managed-postgresql/settings-screenshot.png"

    click_link "Kubernetes"
    click_link "Create Kubernetes Cluster"

    fill_in "Cluster Name", with: "kubernetes-demo"
    find('input[name="cp_nodes"][value="3"]:not([disabled])').trigger("click")
    find('select#worker_nodes option[value="5"]:not([disabled])').select_option
    resize(900)
    screenshot "managed-kubernetes/create-screenshot.png"

    kc_ubid = "kcj4veqfy46a4pcxpmj3dzwzf5"
    kn_ubid = "kna6c1cyxqba5pbmfthkb0xawa"

    k8s_cluster = KubernetesCluster.create(
      name: "kubernetes-demo",
      version: "v1.32",
      cp_node_count: 3,
      location_id: Location::HETZNER_FSN1_ID,
      target_node_size: "standard-2",
      private_subnet_id: Prog::Vnet::SubnetNexus.assemble(
        project.id,
        name: "#{kc_ubid}-subnet",
        location_id: Location::HETZNER_FSN1_ID,
        ipv4_range: Prog::Vnet::SubnetNexus.random_private_ipv4(Location[Location::HETZNER_FSN1_ID], project, 18).to_s
      ).subject.id,
      project_id: project.id
    ) do |kc|
      kc.id = UBID.parse(kc_ubid).to_uuid
    end

    k8s_nodepool = KubernetesNodepool.create(
      name: "kubernetes-demo-np",
      node_count: 5,
      kubernetes_cluster_id: k8s_cluster.id,
      target_node_size: "standard-2"
    ) do |kn|
      kn.id = UBID.parse(kn_ubid).to_uuid
    end

    services_lb = Prog::Vnet::LoadBalancerNexus.assemble(
      k8s_cluster.private_subnet_id,
      name: k8s_cluster.services_load_balancer_name,
      algorithm: "hash_based",
      src_port: 443,
      dst_port: 6443,
      health_check_endpoint: "/",
      health_check_protocol: "tcp",
      stack: LoadBalancer::Stack::IPV4
    ).subject

    # Associate the load balancer with the Kubernetes cluster
    k8s_cluster.update(services_lb_id: services_lb.id)

    ["le9ec", "e12en", "zd3ka"].each do |suffix|
      Prog::Kubernetes::KubernetesNodeNexus.assemble(
        project.id,
        sshable_unix_user: "ubi",
        name: "#{kc_ubid}-#{suffix}",
        location_id: k8s_cluster.location.id,
        private_subnet_id: k8s_cluster.private_subnet_id,
        enable_ip4: true,
        kubernetes_cluster_id: k8s_cluster.id
      ).subject
    end

    ["ubicl", "krn71", "fysgd", "mant4", "m0h11"].each do |suffix|
      Prog::Kubernetes::KubernetesNodeNexus.assemble(
        project.id,
        sshable_unix_user: "ubi",
        name: "#{kn_ubid}-#{suffix}",
        location_id: k8s_cluster.location.id,
        private_subnet_id: k8s_cluster.private_subnet_id,
        enable_ip4: true,
        kubernetes_cluster_id: k8s_cluster.id,
        kubernetes_nodepool_id: k8s_nodepool.id
      ).subject
    end

    Project.define_method(:kubernetes_clusters_dataset) do |*a|
      super(*a).with_extend do
        define_method(:all) { [k8s_cluster] }
        define_method(:first) { |*| k8s_cluster }
      end
    end
    KubernetesCluster.define_method(:display_state) { "running" }
    Vm.define_method(:display_state) { "running" }
    LoadBalancer.define_method(:hostname) { "f9f2x4td37-services.k8s.ubicloud.com" }

    click_link "Kubernetes"
    click_link "kubernetes-demo"
    resize(1400)
    screenshot "managed-kubernetes/overview-screenshot.png"

    resize(800)
    click_link "Users"
    fill_in "email", with: "other@example.com"
    find("input[name=email]").click
    screenshot "security/users-1-screenshot.png"

    account2 = Account.create(email: "other@example.com", name: "Other")
    click_button "Invite"
    page.refresh
    page.evaluate_script("$('#user_policy_#{account2.ubid}').focus()")
    screenshot "security/users-2-screenshot.png"

    click_link "Access Control"
    access_control_path = page.current_path
    screenshot "security/access-control-1-screenshot.png"

    click_link "subject-tags-link"
    screenshot "security/subject-tag-1-screenshot.png"

    fill_in "name", with: "System-Admins"
    click_button "Create"
    fill_in "name", with: "Network-Admins"
    click_button "Create"
    page.refresh
    fill_in "name", with: "Database-Admins"
    page.evaluate_script("$('#name').focus()")
    screenshot "security/subject-tag-2-screenshot.png"
    click_button "Create"

    resize(1020)
    click_link "#{SubjectTag[name: "System-Admins"].ubid}-edit"
    screenshot "security/subject-tag-3-screenshot.png"

    resize(1140)
    check "add[]-#{account2.ubid}-0"
    click_button "Add Members"
    page.refresh
    screenshot "security/subject-tag-4-screenshot.png"

    resize(800)
    visit access_control_path
    click_link "action-tags-link"
    fill_in "name", with: "Networking"
    click_button "Create"
    page.refresh
    screenshot "security/action-tag-1-screenshot.png"

    resize(1140)
    click_link "Manage"
    check "add[]-tazzzzzzzz021gzzzz0fw0a110-0"
    check "add[]-tazzzzzzzz021gzzzz01b0a111-0"
    check "add[]-tazzzzzzzz021gzzzz0ps0a111-0"
    screenshot "security/action-tag-2-screenshot.png"

    resize(800)
    click_button "Add Members"
    page.refresh
    screenshot "security/action-tag-3-screenshot.png"

    resize(850)
    visit access_control_path
    3.times do
      click_button "New Access Control Entry"
    end
    screenshot "security/access-control-2-screenshot.png"

    select "System-Admins", from: "ace-select-1-0"
    select "Vm:all", from: "ace-select-1-1"

    select "Network-Admins", from: "ace-select-2-0"
    select "Networking", from: "ace-select-2-1"

    select "Database-Admins", from: "ace-select-3-0"
    select "Postgres:all", from: "ace-select-3-1"
    screenshot "security/access-control-3-screenshot.png"

    resize(1800)
    click_link "Compute"
    click_link "Create Virtual Machine"
    screenshot "quick-start/managed-services-2-screenshot.png"
  end
end

DB.transaction(rollback: :always, auto_savepoint: true) do |conn|
  DB.temporarily_release_connection(conn) do
    RegenScreenshots.new.call
  end
end

unless RegenScreenshots::SCREENSHOTS.empty?
  warn "Missing screenshots:", RegenScreenshots::SCREENSHOTS.keys.sort
  exit(1)
end