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
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
//! # PyO3 Asyncio Testing Utilities
//!
//! This module provides some utilities for parsing test arguments as well as running and filtering
//! a sequence of tests.
//!
//! As mentioned [here](crate#pythons-event-loop), PyO3 Asyncio tests cannot use the default test
//! harness since it doesn't allow Python to gain control over the main thread. Instead, we have to
//! provide our own test harness in order to create integration tests.
//!
//! Running `pyo3-asyncio` code in doc tests _is_ supported however since each doc test has its own
//! `main` function. When writing doc tests, you may use the
//! [`#[pyo3_asyncio::async_std::main]`](crate::async_std::main) or
//! [`#[pyo3_asyncio::tokio::main]`](crate::tokio::main) macros on the test's main function to run
//! your test.
//!
//! If you don't want to write doc tests, you're unfortunately stuck with integration tests since
//! lib tests do not offer the same level of flexibility for the `main` fn. That being said,
//! overriding the default test harness can be quite different from what you're used to doing for
//! integration tests, so these next sections will walk you through this process.
//!
//! ## Main Test File
//! First, we need to create the test's main file. Although these tests are considered integration
//! tests, we cannot put them in the `tests` directory since that is a special directory owned by
//! Cargo. Instead, we put our tests in a `pytests` directory.
//!
//! > The name `pytests` is just a convention. You can name this folder anything you want in your own
//! > projects.
//!
//! We'll also want to provide the test's main function. Most of the functionality that the test harness needs is packed in the [`pyo3_asyncio::testing::main`](https://docs.rs/pyo3-asyncio/latest/pyo3_asyncio/testing/fn.main.html) function. This function will parse the test's CLI arguments, collect and pass the functions marked with [`#[pyo3_asyncio::async_std::test]`](https://docs.rs/pyo3-asyncio/latest/pyo3_asyncio/async_std/attr.test.html) or [`#[pyo3_asyncio::tokio::test]`](https://docs.rs/pyo3-asyncio/latest/pyo3_asyncio/tokio/attr.test.html) and pass them into the test harness for running and filtering.
//!
//! `pytests/test_example.rs` for the `tokio` runtime:
//! ```rust
//! # #[cfg(all(feature = "tokio-runtime", feature = "attributes"))]
//! #[pyo3_asyncio::tokio::main]
//! async fn main() -> pyo3::PyResult<()> {
//!     pyo3_asyncio::testing::main().await
//! }
//! # #[cfg(not(all(feature = "tokio-runtime", feature = "attributes")))]
//! # fn main() {}
//! ```
//!
//! `pytests/test_example.rs` for the `async-std` runtime:
//! ```rust
//! # #[cfg(all(feature = "async-std-runtime", feature = "attributes"))]
//! #[pyo3_asyncio::async_std::main]
//! async fn main() -> pyo3::PyResult<()> {
//!     pyo3_asyncio::testing::main().await
//! }
//! # #[cfg(not(all(feature = "async-std-runtime", feature = "attributes")))]
//! # fn main() {}
//! ```
//!
//! ## Cargo Configuration
//! Next, we need to add our test file to the Cargo manifest by adding the following section to the
//! `Cargo.toml`
//!
//! ```toml
//! [[test]]
//! name = "test_example"
//! path = "pytests/test_example.rs"
//! harness = false
//! ```
//!
//! Also add the `testing` and `attributes` features to the `pyo3-asyncio` dependency and select your preferred runtime:
//!
//! ```toml
//! pyo3-asyncio = { version = "0.13", features = ["testing", "attributes", "async-std-runtime"] }
//! ```
//!
//! At this point, you should be able to run the test via `cargo test`
//!
//! ### Adding Tests to the PyO3 Asyncio Test Harness
//!
//! We can add tests anywhere in the test crate with the runtime's corresponding `#[test]` attribute:
//!
//! For `async-std` use the [`pyo3_asyncio::async_std::test`](https://docs.rs/pyo3-asyncio/latest/pyo3_asyncio/async_std/attr.test.html) attribute:
//! ```rust
//! # #[cfg(all(feature = "async-std-runtime", feature = "attributes"))]
//! mod tests {
//!     use std::{time::Duration, thread};
//!
//!     use pyo3::prelude::*;
//!
//!     // tests can be async
//!     #[pyo3_asyncio::async_std::test]
//!     async fn test_async_sleep() -> PyResult<()> {
//!         async_std::task::sleep(Duration::from_secs(1)).await;
//!         Ok(())
//!     }
//!
//!     // they can also be synchronous
//!     #[pyo3_asyncio::async_std::test]
//!     fn test_blocking_sleep() -> PyResult<()> {
//!         thread::sleep(Duration::from_secs(1));
//!         Ok(())
//!     }
//! }
//!
//! # #[cfg(all(feature = "async-std-runtime", feature = "attributes"))]
//! #[pyo3_asyncio::async_std::main]
//! async fn main() -> pyo3::PyResult<()> {
//!     pyo3_asyncio::testing::main().await
//! }
//! # #[cfg(not(all(feature = "async-std-runtime", feature = "attributes")))]
//! # fn main() {}
//! ```
//!
//! For `tokio` use the [`pyo3_asyncio::tokio::test`](https://docs.rs/pyo3-asyncio/latest/pyo3_asyncio/tokio/attr.test.html) attribute:
//! ```rust
//! # #[cfg(all(feature = "tokio-runtime", feature = "attributes"))]
//! mod tests {
//!     use std::{time::Duration, thread};
//!
//!     use pyo3::prelude::*;
//!
//!     // tests can be async
//!     #[pyo3_asyncio::tokio::test]
//!     async fn test_async_sleep() -> PyResult<()> {
//!         tokio::time::sleep(Duration::from_secs(1)).await;
//!         Ok(())
//!     }
//!
//!     // they can also be synchronous
//!     #[pyo3_asyncio::tokio::test]
//!     fn test_blocking_sleep() -> PyResult<()> {
//!         thread::sleep(Duration::from_secs(1));
//!         Ok(())
//!     }
//! }
//!
//! # #[cfg(all(feature = "tokio-runtime", feature = "attributes"))]
//! #[pyo3_asyncio::tokio::main]
//! async fn main() -> pyo3::PyResult<()> {
//!     pyo3_asyncio::testing::main().await
//! }
//! # #[cfg(not(all(feature = "tokio-runtime", feature = "attributes")))]
//! # fn main() {}
//! ```
//!
//! ## Lib Tests
//!
//! Unfortunately, as we mentioned at the beginning, these utilities will only run in integration
//! tests and doc tests. Running lib tests are out of the question since we need control over the
//! main function. You can however perform compilation checks for lib tests. This is much more
//! useful in doc tests than it is for lib tests, but the option is there if you want it.
//!
//! `my-crate/src/lib.rs`
//! ```
//! # #[cfg(all(
//! #     any(feature = "async-std-runtime", feature = "tokio-runtime"),
//! #     feature = "attributes"
//! # ))]
//! mod tests {
//!     use pyo3::prelude::*;
//!
//! #   #[cfg(feature = "async-std-runtime")]
//!     #[pyo3_asyncio::async_std::test]
//!     async fn test_async_std_async_test_compiles() -> PyResult<()> {
//!         Ok(())
//!     }
//! #   #[cfg(feature = "async-std-runtime")]
//!     #[pyo3_asyncio::async_std::test]
//!     fn test_async_std_sync_test_compiles() -> PyResult<()> {
//!         Ok(())
//!     }
//!
//! #   #[cfg(feature = "tokio-runtime")]
//!     #[pyo3_asyncio::tokio::test]
//!     async fn test_tokio_async_test_compiles() -> PyResult<()> {
//!         Ok(())
//!     }
//! #   #[cfg(feature = "tokio-runtime")]
//!     #[pyo3_asyncio::tokio::test]
//!     fn test_tokio_sync_test_compiles() -> PyResult<()> {
//!         Ok(())
//!     }
//! }
//!
//! # fn main() {}
//! ```

use std::{future::Future, pin::Pin};

use clap::{Arg, Command};
use futures::stream::{self, StreamExt};
use pyo3::prelude::*;

/// Args that should be provided to the test program
///
/// These args are meant to mirror the default test harness's args.
/// > Currently only `--filter` is supported.
pub struct Args {
    filter: Option<String>,
}

impl Default for Args {
    fn default() -> Self {
        Self { filter: None }
    }
}

/// Parse the test args from the command line
///
/// This should be called at the start of your test harness to give the CLI some
/// control over how our tests are run.
///
/// Ideally, we should mirror the default test harness's arguments exactly, but
/// for the sake of simplicity, only filtering is supported for now. If you want
/// more features, feel free to request them
/// [here](https://github.com/awestlake87/pyo3-asyncio/issues).
///
/// # Examples
///
/// Running the following function:
/// ```
/// # use pyo3_asyncio::testing::parse_args;
/// let args = parse_args();
/// ```
///
/// Produces the following usage string:
///
/// ```bash
/// Pyo3 Asyncio Test Suite
/// USAGE:
/// test_example [TESTNAME]
///
/// FLAGS:
/// -h, --help       Prints help information
/// -V, --version    Prints version information
///
/// ARGS:
/// <TESTNAME>    If specified, only run tests containing this string in their names
/// ```
pub fn parse_args() -> Args {
    let matches = Command::new("PyO3 Asyncio Test Suite")
        .arg(
            Arg::new("TESTNAME")
                .help("If specified, only run tests containing this string in their names"),
        )
        .get_matches();

    Args {
        filter: matches
            .get_one::<String>("TESTNAME")
            .map(|name| name.clone()),
    }
}

type TestFn = dyn Fn() -> Pin<Box<dyn Future<Output = PyResult<()>> + Send>> + Send + Sync;

/// The structure used by the `#[test]` macros to provide a test to the `pyo3-asyncio` test harness.
#[derive(Clone)]
pub struct Test {
    /// The fully qualified name of the test
    pub name: &'static str,
    /// The function used to create the task that runs the test.
    pub test_fn: &'static TestFn,
}

impl Test {
    /// Create the task that runs the test
    pub fn task(
        &self,
    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = pyo3::PyResult<()>> + Send>> {
        (self.test_fn)()
    }
}

inventory::collect!(Test);

/// Run a sequence of tests while applying any necessary filtering from the `Args`
pub async fn test_harness(tests: Vec<Test>, args: Args) -> PyResult<()> {
    stream::iter(tests)
        .for_each_concurrent(Some(4), |test| {
            let mut ignore = false;

            if let Some(filter) = args.filter.as_ref() {
                if !test.name.contains(filter) {
                    ignore = true;
                }
            }

            async move {
                if !ignore {
                    test.task().await.unwrap();

                    println!("test {} ... ok", test.name);
                }
            }
        })
        .await;

    Ok(())
}

/// Parses test arguments and passes the tests to the `pyo3-asyncio` test harness
///
/// This function collects the test structures from the `inventory` boilerplate and forwards them to
/// the test harness.
///
/// # Examples
///
/// ```
/// # #[cfg(all(feature = "async-std-runtime", feature = "attributes"))]
/// use pyo3::prelude::*;
///
/// # #[cfg(all(feature = "async-std-runtime", feature = "attributes"))]
/// #[pyo3_asyncio::async_std::main]
/// async fn main() -> PyResult<()> {
///     pyo3_asyncio::testing::main().await
/// }
/// # #[cfg(not(all(feature = "async-std-runtime", feature = "attributes")))]
/// # fn main() { }
/// ```
pub async fn main() -> PyResult<()> {
    let args = parse_args();

    test_harness(
        inventory::iter::<Test>().map(|test| test.clone()).collect(),
        args,
    )
    .await
}

#[cfg(test)]
#[cfg(all(
    feature = "testing",
    feature = "attributes",
    any(feature = "async-std-runtime", feature = "tokio-runtime")
))]
mod tests {
    use pyo3::prelude::*;

    use crate as pyo3_asyncio;

    #[cfg(feature = "async-std-runtime")]
    #[pyo3_asyncio::async_std::test]
    async fn test_async_std_async_test_compiles() -> PyResult<()> {
        Ok(())
    }
    #[cfg(feature = "async-std-runtime")]
    #[pyo3_asyncio::async_std::test]
    fn test_async_std_sync_test_compiles() -> PyResult<()> {
        Ok(())
    }

    #[cfg(feature = "tokio-runtime")]
    #[pyo3_asyncio::tokio::test]
    async fn test_tokio_async_test_compiles() -> PyResult<()> {
        Ok(())
    }
    #[cfg(feature = "tokio-runtime")]
    #[pyo3_asyncio::tokio::test]
    fn test_tokio_sync_test_compiles() -> PyResult<()> {
        Ok(())
    }
}