1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141

//! Additional combinators for testing async writers.

use futures_io::AsyncWrite;

pub use super::limited::Limited;
pub use crate::assert_unmoved::AssertUnmoved;
pub use crate::interleave_pending::InterleavePending;
pub use crate::track_closed::TrackClosed;

/// Additional combinators for testing async writers.
pub trait AsyncWriteTestExt: AsyncWrite {
    /// Asserts that the given is not moved after being polled.
    ///
    /// A check for movement is performed each time the writer is polled
    /// and when `Drop` is called.
    ///
    /// Aside from keeping track of the location at which the writer was first
    /// polled and providing assertions, this writer adds no runtime behavior
    /// and simply delegates to the child writer.
    fn assert_unmoved_write(self) -> AssertUnmoved<Self>
    where
        Self: Sized,
    {
        AssertUnmoved::new(self)
    }

    /// Introduces an extra [`Poll::Pending`](futures_core::task::Poll::Pending)
    /// in between each operation on the writer.
    ///
    /// # Examples
    ///
    /// ```
    /// use futures::task::Poll;
    /// use futures::io::{AsyncWrite, Cursor};
    /// use futures_test::task::noop_context;
    /// use futures_test::io::AsyncWriteTestExt;
    /// use futures::pin_mut;
    ///
    /// let writer = Cursor::new(vec![0u8; 4].into_boxed_slice()).interleave_pending_write();
    /// pin_mut!(writer);
    ///
    /// let mut cx = noop_context();
    ///
    /// assert_eq!(writer.as_mut().poll_write(&mut cx, &[1, 2])?, Poll::Pending);
    /// assert_eq!(writer.as_mut().poll_write(&mut cx, &[1, 2])?, Poll::Ready(2));
    /// assert_eq!(&writer.get_ref().get_ref()[..], [1, 2, 0, 0]);
    /// assert_eq!(writer.as_mut().poll_write(&mut cx, &[3, 4])?, Poll::Pending);
    /// assert_eq!(writer.as_mut().poll_write(&mut cx, &[3, 4])?, Poll::Ready(2));
    /// assert_eq!(&writer.get_ref().get_ref()[..], [1, 2, 3, 4]);
    /// assert_eq!(writer.as_mut().poll_write(&mut cx, &[5, 6])?, Poll::Pending);
    /// assert_eq!(writer.as_mut().poll_write(&mut cx, &[5, 6])?, Poll::Ready(0));
    ///
    /// assert_eq!(writer.as_mut().poll_flush(&mut cx)?, Poll::Pending);
    /// assert_eq!(writer.as_mut().poll_flush(&mut cx)?, Poll::Ready(()));
    ///
    /// assert_eq!(writer.as_mut().poll_close(&mut cx)?, Poll::Pending);
    /// assert_eq!(writer.as_mut().poll_close(&mut cx)?, Poll::Ready(()));
    ///
    /// # Ok::<(), std::io::Error>(())
    /// ```
    fn interleave_pending_write(self) -> InterleavePending<Self>
    where
        Self: Sized,
    {
        InterleavePending::new(self)
    }

    /// Limit the number of bytes allowed to be written on each call to `poll_write`.
    ///
    /// # Examples
    ///
    /// ```
    /// use futures::task::Poll;
    /// use futures::io::{AsyncWrite, Cursor};
    /// use futures_test::task::noop_context;
    /// use futures_test::io::AsyncWriteTestExt;
    /// use futures::pin_mut;
    ///
    /// let writer = Cursor::new(vec![0u8; 4].into_boxed_slice()).limited_write(2);
    /// pin_mut!(writer);
    ///
    /// let mut cx = noop_context();
    ///
    /// assert_eq!(writer.as_mut().poll_write(&mut cx, &[1, 2])?, Poll::Ready(2));
    /// assert_eq!(&writer.get_ref().get_ref()[..], [1, 2, 0, 0]);
    /// assert_eq!(writer.as_mut().poll_write(&mut cx, &[3])?, Poll::Ready(1));
    /// assert_eq!(&writer.get_ref().get_ref()[..], [1, 2, 3, 0]);
    /// assert_eq!(writer.as_mut().poll_write(&mut cx, &[4, 5])?, Poll::Ready(1));
    /// assert_eq!(&writer.get_ref().get_ref()[..], [1, 2, 3, 4]);
    /// assert_eq!(writer.as_mut().poll_write(&mut cx, &[5])?, Poll::Ready(0));
    ///
    /// # Ok::<(), std::io::Error>(())
    /// ```
    fn limited_write(self, limit: usize) -> Limited<Self>
    where
        Self: Sized,
    {
        Limited::new(self, limit)
    }

    /// Track whether this stream has been closed and errors if it is used after closing.
    ///
    /// # Examples
    ///
    /// ```
    /// # futures::executor::block_on(async {
    /// use futures::io::{AsyncWriteExt, Cursor};
    /// use futures_test::io::AsyncWriteTestExt;
    ///
    /// let mut writer = Cursor::new(vec![0u8; 4]).track_closed();
    ///
    /// writer.write_all(&[1, 2]).await?;
    /// assert!(!writer.is_closed());
    /// writer.close().await?;
    /// assert!(writer.is_closed());
    ///
    /// # Ok::<(), std::io::Error>(()) })?;
    /// # Ok::<(), std::io::Error>(())
    /// ```
    ///
    /// ```
    /// # futures::executor::block_on(async {
    /// use futures::io::{AsyncWriteExt, Cursor};
    /// use futures_test::io::AsyncWriteTestExt;
    ///
    /// let mut writer = Cursor::new(vec![0u8; 4]).track_closed();
    ///
    /// writer.close().await?;
    /// assert!(writer.write_all(&[1, 2]).await.is_err());
    /// # Ok::<(), std::io::Error>(()) })?;
    /// # Ok::<(), std::io::Error>(())
    /// ```
    fn track_closed(self) -> TrackClosed<Self>
    where
        Self: Sized,
    {
        TrackClosed::new(self)
    }
}

impl<W> AsyncWriteTestExt for W where W: AsyncWrite {}