3use support::mock_jev::{MockJev, Reply};
4use support::{jev_instance, noul};
5use tokio_postgres::error::SqlState;

7#[tokio::test(flavor = "multi_thread")]
8async fn only_granted_roles_may_spend() {
9    let mock = MockJev::start(|_| Reply::json(200, noul(0.95))).await;
10    let (_pg, client) = jev_instance(&mock).await;
11    client
12        .batch_execute("CREATE ROLE app; GRANT SELECT ON tickets TO app; SET ROLE app;")
13        .await
14        .unwrap();
15
16    let err = client
17        .query_one("SELECT jev_prob(t, 'Urgent?') FROM tickets t", &[])
18        .await
19        .expect_err("not granted");
20    assert_eq!(err.as_db_error().unwrap().code(), &SqlState::INSUFFICIENT_PRIVILEGE);
21    assert!(mock.requests().is_empty());
22
23    client
24        .batch_execute("RESET ROLE; GRANT EXECUTE ON FUNCTION jev_prob(anyelement, text) TO app; SET ROLE app;")
25        .await
26        .unwrap();
27    let row = client.query_one("SELECT jev_prob(t, 'Urgent?') FROM tickets t", &[]).await.expect("granted");
28    assert_eq!(row.get::<_, f64>(0), 0.95);
29}

33#[tokio::test(flavor = "multi_thread")]
34async fn no_function_is_executable_by_public() {
35    let mock = MockJev::start(|_| Reply::json(200, noul(0.95))).await;
36    let (_pg, client) = jev_instance(&mock).await;
37    let functions: Vec<(String, bool)> = client
38        .query(
39            "SELECT p.oid::regprocedure::text, has_function_privilege('public', p.oid, 'EXECUTE')
40             FROM pg_proc p
41             JOIN pg_depend d ON d.classid = 'pg_proc'::regclass AND d.objid = p.oid AND d.deptype = 'e'
42             JOIN pg_extension e ON e.oid = d.refobjid
43             WHERE e.extname = 'postjevsql'
44             ORDER BY 1",
45            &[],
46        )
47        .await
48        .unwrap()
49        .iter()
50        .map(|r| (r.get(0), r.get(1)))
51        .collect();
52    // The query must see the functions, or an empty result proves nothing.
53    let names: Vec<&str> = functions.iter().map(|(f, _)| f.as_str()).collect();
54    for expected in ["jev(anyelement,text)", "jev_prob(anyelement,text)", "jev_stats()"] {
55        assert!(names.contains(&expected), "{expected} missing from {names:?}");
56    }
57    assert!(names.iter().all(|f| f.starts_with("jev")), "{names:?}");
58    let open: Vec<&str> = functions.iter().filter(|(_, open)| *open).map(|(f, _)| f.as_str()).collect();
59    assert_eq!(open, Vec::<&str>::new());
60}