/lib/k8s-test-framework/src/reader.rs

https://github.com/timberio/vector · Rust · 123 lines · 79 code · 21 blank · 23 comment · 6 complexity · dada792159f01b324ad3b14cd30b2953 MD5 · raw file

  1. //! Read process output.
  2. use std::process::{ExitStatus, Stdio};
  3. use tokio::io::{AsyncBufReadExt, BufReader};
  4. use tokio::process::{Child, ChildStdout, Command};
  5. /// Keeps track of the command invocation, proving the interface to
  6. /// read the output and send a termination signal.
  7. #[derive(Debug)]
  8. pub struct Reader {
  9. child: Child,
  10. reader: BufReader<ChildStdout>,
  11. }
  12. impl Reader {
  13. /// Spawn a command and provide a [`Reader`].
  14. pub fn spawn(mut command: Command) -> std::io::Result<Self> {
  15. Self::prepare_stdout(&mut command);
  16. let child = command.spawn()?;
  17. Ok(Self::new(child))
  18. }
  19. fn prepare_stdout(command: &mut Command) {
  20. command.stdout(Stdio::piped());
  21. }
  22. fn new(mut child: Child) -> Self {
  23. let stdout = child.stdout.take().unwrap();
  24. let reader = BufReader::new(stdout);
  25. Reader { child, reader }
  26. }
  27. /// Wait for the `kubectl logs` process to exit and return the exit code.
  28. pub async fn wait(&mut self) -> std::io::Result<ExitStatus> {
  29. (&mut self.child).await
  30. }
  31. /// Send a termination signal to the `kubectl logs` process.
  32. pub fn kill(&mut self) -> std::io::Result<()> {
  33. self.child.kill()
  34. }
  35. /// Read one line from the stdout of the `kubectl logs` process.
  36. pub async fn read_line(&mut self) -> Option<String> {
  37. let mut s = String::new();
  38. let result = self.reader.read_line(&mut s).await;
  39. match result {
  40. Ok(0) => None,
  41. Ok(_) => Some(s),
  42. Err(err) => panic!(err),
  43. }
  44. }
  45. }
  46. #[cfg(test)]
  47. mod tests {
  48. use super::*;
  49. async fn collect(reader: &mut Reader) -> Vec<String> {
  50. let mut list = Vec::new();
  51. while let Some(line) = reader.read_line().await {
  52. list.push(line)
  53. }
  54. list
  55. }
  56. #[tokio::test]
  57. async fn test_reader_finite() {
  58. let mut command = Command::new("echo");
  59. command.arg("test");
  60. let mut reader = Reader::spawn(command).expect("unable to spawn");
  61. // Collect all line, expect stream to finish.
  62. let lines = collect(&mut reader).await;
  63. // Assert we got all the lines we expected.
  64. assert_eq!(lines, vec!["test\n".to_owned()]);
  65. // Ensure wait doesn't fail, and that we exit status is success.
  66. let exit_status = reader.wait().await.expect("wait failed");
  67. assert!(exit_status.success());
  68. }
  69. #[tokio::test]
  70. async fn test_reader_infinite() {
  71. let mut command = Command::new("bash");
  72. command.arg("-c");
  73. command.arg(r#"NUM=0; while true; do echo "Line $NUM"; NUM=$((NUM+=1)); sleep 0.01; done"#);
  74. let mut reader = Reader::spawn(command).expect("unable to spawn");
  75. // Read the lines and at some point ask the command we're reading from
  76. // to stop.
  77. let mut expected_num = 0;
  78. while let Some(line) = reader.read_line().await {
  79. // Assert we're getting expected lines.
  80. assert_eq!(line, format!("Line {}\n", expected_num));
  81. // On line 100 issue a `kill` to stop the infinite stream.
  82. if expected_num == 100 {
  83. reader.kill().expect("process already stopped")
  84. }
  85. // If we are past 200 it means we issued `kill` at 100 and it wasn't
  86. // effective. This is problem, fail the test.
  87. // We don't to this immediately after `kill` to allow for some
  88. // potential race condition. That kind of race is not just ok, but
  89. // is desirable in the real-life usage to read-up the whole stdout
  90. // buffer.
  91. if expected_num > 200 {
  92. panic!("Went too far without stop being effective");
  93. }
  94. // Bump the expected num for the next iteration.
  95. expected_num += 1;
  96. }
  97. // Ensure wait doesn't fail. We killed the process, so expect
  98. // a non-success exit code.
  99. let exit_status = reader.wait().await.expect("wait failed");
  100. assert!(!exit_status.success());
  101. }
  102. }