Repository: firecracker-microvm/firecracker
File: tests/integration_tests/functional/test_vsock.py (Python)

Code before:
    test_vm.start()
    test_vm.wait_for_up()

    # Generate the random data blob file.
    blob_path, blob_hash = make_blob(test_fc_session_root_path)
    vm_blob_path = "/tmp/vsock/test.blob"

    # Set up a tmpfs drive on the guest, so we can copy the blob there.
    # Guest-initiated connections (echo workers) will use this blob.
    _copy_vsock_data_to_guest(test_vm.ssh, blob_path, vm_blob_path, bin_vsock_path)

    host_socket_path = os.path.join(
        test_vm.path, f"{VSOCK_UDS_PATH}_{ECHO_SERVER_PORT}"
    )
    host_socat_commmand = [
        "socat",
        "-dddd",
        f"UNIX-LISTEN:{host_socket_path},fork",
        "STDOUT",
    ]
    host_socat = subprocess.Popen(
        host_socat_commmand, stdout=subprocess.PIPE, stderr=subprocess.PIPE
    )

    # Give some time for host socat to create socket
    time.sleep(0.5)
    assert Path(host_socket_path).exists()
    test_vm.create_jailed_resource(host_socket_path)

    # Create a socat process in the guest which will connect to the host socat
    guest_socat_commmand = f"tmux new -d 'socat - vsock-connect:2:{ECHO_SERVER_PORT}'"
    test_vm.ssh.run(guest_socat_commmand)

    # socat should be running in the guest now
    code, _, _ = test_vm.ssh.run("pidof socat")
    assert code == 0

    # Create snapshot.
    snapshot = test_vm.snapshot_full()
    test_vm.resume()

    # After `create_snapshot` + 'restore' calls, connection should be dropped
    code, _, _ = test_vm.ssh.run("pidof socat")
    assert code == 1

    # Kill host socat as it is not useful anymore
    host_socat.kill()
    host_socat.communicate()

    # Terminate VM.
    metrics = test_vm.flush_metrics()

Review comment (suggestion):
```suggestion
    test_vm.ssh.check_output("pidof socat")
```

Code after:

    host_socket_path = os.path.join(
        test_vm.path, f"{VSOCK_UDS_PATH}_{ECHO_SERVER_PORT}"
    )
    host_socat_commmand = [
        "socat",
        "-dddd",
        f"UNIX-LISTEN:{host_socket_path},fork",
        "STDOUT",
    ]
    host_socat = subprocess.Popen(
        host_socat_commmand, stdout=subprocess.PIPE, stderr=subprocess.PIPE
    )

    # Give some time for host socat to create socket
    time.sleep(0.5)
    assert Path(host_socket_path).exists()
    test_vm.create_jailed_resource(host_socket_path)

    # Create a socat process in the guest which will connect to the host socat
    guest_socat_commmand = f"tmux new -d 'socat - vsock-connect:2:{ECHO_SERVER_PORT}'"
    test_vm.ssh.run(guest_socat_commmand)

    # socat should be running in the guest now
    code, _, _ = test_vm.ssh.run("pidof socat")
    assert code == 0

    # Create snapshot.
    snapshot = test_vm.snapshot_full()
    test_vm.resume()

    # After `create_snapshot` + 'restore' calls, connection should be dropped
    code, _, _ = test_vm.ssh.run("pidof socat")
    assert code == 1

    # Kill host socat as it is not useful anymore
    host_socat.kill()
    host_socat.communicate()

    # Terminate VM.
    test_vm.kill()

    # Load snapshot.
    vm2 = microvm_factory.build()
    vm2.spawn()
    vm2.restore_from_snapshot(snapshot, resume=True)
    vm2.wait_for_up()

    # After snap restore all vsock connections should be
    # dropped. This means guest socat should exit same way
    # as it did after snapshot was taken.