test: add a/b test checker

codegen/Cargo.toml

@@ -18,3 +18,5 @@ clap = { version = "3.2.16", features = ["derive"] }
 phf_codegen = "0.11.1"
 once_cell = "1.12.0"
 fancy-regex = "0.10.0"
+indicatif = "0.17.0"
+num_enum = "0.5.7"

codegen/src/abtest.rs

@@ -0,0 +1,112 @@
+use anyhow::{bail, Result};
+use futures::{stream, StreamExt};
+use indicatif::{ProgressBar, ProgressStyle};
+use num_enum::TryFromPrimitive;
+use rustypipe::client::{ClientType, RustyPipe, YTContext};
+use serde::{Deserialize, Serialize};
+
+#[derive(
+    Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, TryFromPrimitive, Serialize, Deserialize,
+)]
+#[repr(u16)]
+pub enum ABTest {
+    AttributedTextDescription = 1,
+    ThreeTabChannelLayout = 2,
+}
+
+const N_TESTS: u16 = 2;
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ABTestRes {
+    id: u16,
+    name: ABTest,
+    tests: usize,
+    occurrences: usize,
+}
+
+#[derive(Debug, Serialize)]
+struct QVideo<'a> {
+    context: YTContext<'a>,
+    video_id: &'a str,
+    content_check_ok: bool,
+    racy_check_ok: bool,
+}
+
+pub async fn run_test(ab: ABTest, n: usize, concurrency: usize) -> usize {
+    eprintln!("🧪 A/B test #{}: {:?}", ab as u16, ab);
+
+    let rp = RustyPipe::new();
+    let pb = ProgressBar::new(n as u64);
+    pb.set_style(
+        ProgressStyle::with_template(
+            "{msg} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {pos}/{len}",
+        )
+        .unwrap(),
+    );
+    // let mut count = 0;
+
+    let results = stream::iter(0..n)
+        .map(|_| {
+            let rp = rp.clone();
+            let pb = pb.clone();
+            async move {
+                let is_present = match ab {
+                    ABTest::AttributedTextDescription => attributed_text_description(&rp).await,
+                    ABTest::ThreeTabChannelLayout => three_tab_channel_layout(&rp).await,
+                }
+                .unwrap();
+                pb.inc(1);
+                is_present
+            }
+        })
+        .buffer_unordered(concurrency)
+        .collect::<Vec<_>>()
+        .await;
+
+    let count = results.iter().filter(|x| **x).count();
+
+    count
+}
+
+pub async fn run_all_tests(n: usize, concurrency: usize) -> Vec<ABTestRes> {
+    let mut results = Vec::new();
+
+    for id in 1..=N_TESTS {
+        let ab = ABTest::try_from(id).unwrap();
+        let occurrences = run_test(ab, n, concurrency).await;
+        results.push(ABTestRes {
+            id,
+            name: ab,
+            tests: n,
+            occurrences,
+        });
+    }
+    results
+}
+
+pub async fn attributed_text_description(rp: &RustyPipe) -> Result<bool> {
+    let query = rp.query();
+    let context = query.get_context(ClientType::Desktop, true, None).await;
+    let q = QVideo {
+        context,
+        video_id: "ZeerrnuLi5E",
+        content_check_ok: false,
+        racy_check_ok: false,
+    };
+    let response_txt = query.raw(ClientType::Desktop, "next", &q).await.unwrap();
+
+    if !response_txt.contains("\"Black Mamba\"") {
+        bail!("invalid response data");
+    }
+
+    Ok(response_txt.contains("\"attributedDescription\""))
+}
+
+pub async fn three_tab_channel_layout(rp: &RustyPipe) -> Result<bool> {
+    let channel = rp
+        .query()
+        .channel_videos("UCR-DXc1voovS8nhAvccRZhg")
+        .await
+        .unwrap();
+    Ok(channel.has_live || channel.has_shorts)
+}

codegen/src/main.rs

@@ -1,3 +1,4 @@
+mod abtest;
 mod collect_album_types;
 mod collect_large_numbers;
 mod collect_playlist_dates;
@@ -31,6 +32,12 @@ enum Commands {
     GenLocales,
     GenDict,
     DownloadTestfiles,
+    AbTest {
+        #[clap(value_parser)]
+        id: Option<u16>,
+        #[clap(short, default_value = "100")]
+        n: usize,
+    },
 }
 
 #[tokio::main]
@@ -62,5 +69,18 @@ async fn main() {
         Commands::DownloadTestfiles => {
             download_testfiles::download_testfiles(&cli.project_root).await
         }
+        Commands::AbTest { id, n } => {
+            match id {
+                Some(id) => {
+                    let ab = abtest::ABTest::try_from(id).expect("invalid A/B test id");
+                    let res = abtest::run_test(ab, n, cli.concurrency).await;
+                    eprintln!("{} occurences", res);
+                }
+                None => {
+                    let res = abtest::run_all_tests(n, cli.concurrency).await;
+                    println!("{}", serde_json::to_string(&res).unwrap())
+                }
+            };
+        }
     };
 }