make not totally broken (GPT-6)

arbtt-sway/src/main.rs

@@ -1,5 +1,8 @@
 // based on https://github.com/josephdunn/idlers/blob/main/src/main.rs
 
+#[cfg(test)]
+mod tests;
+
 use std::process::{Command, Stdio};
 use std::time::{Instant, Duration};
 use std::sync::{Mutex, Arc};
@@ -7,7 +10,7 @@ use std::io::Write;
 use std::thread;
 use anyhow::Result;
 use chrono::{DateTime, Utc};
-use swayipc::{EventType, Node, NodeType, Workspace};
+use swayipc::{EventType, Node, NodeType};
 use wayland_client::{Connection, Dispatch, QueueHandle, delegate_noop, event_created_child, globals::{GlobalListContents, registry_queue_init}, protocol::{wl_registry, wl_seat}};
 use wayland_protocols::ext::idle_notify::v1::client::{ext_idle_notification_v1::{self, ExtIdleNotificationV1}, ext_idle_notifier_v1::ExtIdleNotifierV1};
 use smol_str::SmolStr;
@@ -26,7 +29,7 @@ impl Dispatch<ExtIdleNotificationV1, ()> for WrState {
         state: &mut Self,
         _proxy: &ExtIdleNotificationV1,
         event: ext_idle_notification_v1::Event,
-        data: &(),
+        _data: &(),
         _conn: &Connection,
         _qh: &QueueHandle<Self>,
     ) {
@@ -100,12 +103,19 @@ fn process_tree<'a>(node: &'a Node, windows: &mut Vec<ArbttWindow>, focused: &mu
 }
 
 fn sway_thread(wr_state: WrState) -> Result<()> {
-    let mut connection = swayipc::Connection::new()?;
+    watch_sway(wr_state, swayipc::Connection::new()?, swayipc::Connection::new()?)
+}
 
-    let tree = connection.get_tree()?;
+fn watch_sway(
+    wr_state: WrState,
+    mut connection: swayipc::Connection,
+    events: swayipc::Connection,
+) -> Result<()> {
+    let events = events.subscribe([EventType::Window, EventType::Workspace])?;
     let mut windows = vec![];
 
-    let mut read_windows = || {
+    let mut read_windows = || -> Result<()> {
+        let tree = connection.get_tree()?;
         windows.clear();
         let mut focused = SmolStr::new("");
         process_tree(&tree, &mut windows, &mut focused, None);
@@ -113,17 +123,17 @@ fn sway_thread(wr_state: WrState) -> Result<()> {
         //println!("emit tree {:?}", windows);
         std::mem::swap(&mut state.windows, &mut windows);
         state.focused_desktop = focused;
+        Ok(())
     };
 
-    read_windows();
+    read_windows()?;
 
-    let st = connection.subscribe([EventType::Window])?;
-    for _s in st {
-        //println!("reread");
-        read_windows();
+    for event in events {
+        event?;
+        read_windows()?;
     }
 
-    std::process::exit(0);
+    anyhow::bail!("Sway event stream ended")
 }
 
 const INTERVAL: Duration = Duration::from_secs(60);
@@ -170,7 +180,10 @@ fn main() -> Result<()> {
     let wr_state_ = wr_state.clone();
     let wr_state__ = wr_state.clone();
     thread::spawn(|| {
-        sway_thread(wr_state_).unwrap()
+        if let Err(error) = sway_thread(wr_state_) {
+            eprintln!("Sway window tracking failed: {error:#}");
+            std::process::exit(1);
+        }
     });
     thread::spawn(|| {
        arbtt_thread(wr_state__).unwrap()

arbtt-sway/src/tests.rs

@@ -0,0 +1,129 @@
+use super::*;
+use serde_json::{Value, json};
+use std::io::Read;
+use std::os::unix::net::UnixStream;
+
+fn socket_pair() -> (swayipc::Connection, UnixStream) {
+    let (client, server) = UnixStream::pair().unwrap();
+    for socket in [&client, &server] {
+        socket
+            .set_read_timeout(Some(Duration::from_secs(5)))
+            .unwrap();
+        socket
+            .set_write_timeout(Some(Duration::from_secs(5)))
+            .unwrap();
+    }
+    (client.into(), server)
+}
+
+fn receive(socket: &mut UnixStream, expected_type: u32) -> Vec<u8> {
+    let mut header = [0; 14];
+    socket.read_exact(&mut header).unwrap();
+    assert_eq!(&header[..6], b"i3-ipc");
+    assert_eq!(
+        u32::from_ne_bytes(header[10..14].try_into().unwrap()),
+        expected_type
+    );
+    let mut payload = vec![0; u32::from_ne_bytes(header[6..10].try_into().unwrap()) as usize];
+    socket.read_exact(&mut payload).unwrap();
+    payload
+}
+
+fn send(socket: &mut UnixStream, message_type: u32, payload: &Value) {
+    let payload = serde_json::to_vec(payload).unwrap();
+    socket.write_all(b"i3-ipc").unwrap();
+    socket
+        .write_all(&(payload.len() as u32).to_ne_bytes())
+        .unwrap();
+    socket.write_all(&message_type.to_ne_bytes()).unwrap();
+    socket.write_all(&payload).unwrap();
+}
+
+fn node(kind: &str, name: &str, focused: bool, children: Vec<Value>) -> Value {
+    let rect = json!({"x": 0, "y": 0, "width": 100, "height": 100});
+    json!({
+        "id": 1, "name": name, "type": kind, "border": "none",
+        "current_border_width": 0, "layout": "splith", "orientation": "none",
+        "rect": rect, "window_rect": rect, "deco_rect": rect, "geometry": rect,
+        "urgent": false, "focused": focused, "focus": [], "nodes": children,
+        "floating_nodes": [], "sticky": false
+    })
+}
+
+#[test]
+fn refreshes_empty_startup_tree_on_window_and_workspace_events() {
+    let state = WrState(Arc::new(Mutex::new(State {
+        idle_since: None,
+        windows: vec![],
+        focused_desktop: SmolStr::new(""),
+    })));
+    let (queries, mut query_server) = socket_pair();
+    let (events, mut event_server) = socket_pair();
+    let worker_state = state.clone();
+    let worker = thread::spawn(move || watch_sway(worker_state, queries, events));
+
+    let subscriptions: Value = serde_json::from_slice(&receive(&mut event_server, 2)).unwrap();
+    assert_eq!(subscriptions, json!(["window", "workspace"]));
+    // The subscription must complete before the first tree query.
+    query_server.set_nonblocking(true).unwrap();
+    assert_eq!(
+        query_server.read(&mut [0]).unwrap_err().kind(),
+        std::io::ErrorKind::WouldBlock
+    );
+    query_server.set_nonblocking(false).unwrap();
+    send(&mut event_server, 2, &json!({"success": true}));
+
+    receive(&mut query_server, 4);
+    send(&mut query_server, 4, &node("workspace", "1", true, vec![]));
+
+    let opened = node("con", "First title", true, vec![]);
+    send(
+        &mut event_server,
+        0x80000003,
+        &json!({"change": "new", "container": opened}),
+    );
+    receive(&mut query_server, 4);
+    assert!(state.0.lock().unwrap().windows.is_empty());
+    send(
+        &mut query_server,
+        4,
+        &node("workspace", "1", false, vec![opened]),
+    );
+
+    let renamed = node("con", "Updated title", true, vec![]);
+    send(
+        &mut event_server,
+        0x80000003,
+        &json!({"change": "title", "container": renamed}),
+    );
+    receive(&mut query_server, 4);
+    // Receiving the next query proves the previous snapshot has been published.
+    assert_eq!(state.0.lock().unwrap().windows[0].title, "First title");
+    send(
+        &mut query_server,
+        4,
+        &node("workspace", "1", false, vec![renamed]),
+    );
+
+    let empty_workspace = node("workspace", "2", true, vec![]);
+    send(
+        &mut event_server,
+        0x80000000,
+        &json!({"change": "focus", "current": empty_workspace}),
+    );
+    receive(&mut query_server, 4);
+    {
+        let state = state.0.lock().unwrap();
+        assert_eq!(state.windows[0].title, "Updated title");
+        assert!(state.windows[0].active);
+        assert_eq!(state.focused_desktop, "1");
+    }
+    send(&mut query_server, 4, &empty_workspace);
+
+    // A disconnected event socket must return an error, not keep logging stale data.
+    drop(event_server);
+    assert!(worker.join().unwrap().is_err());
+    let state = state.0.lock().unwrap();
+    assert!(state.windows.is_empty());
+    assert_eq!(state.focused_desktop, "2");
+}