From 8ed0fb751cb3d9fc40b3bd12f5d48c0f9377c5f5 Mon Sep 17 00:00:00 2001
From: NexVeridian <NexVeridian@gmail.com>
Date: Wed, 24 Apr 2024 18:01:39 -0700
Subject: [PATCH] ARKB for ARKF

---
 src/util.rs | 46 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/src/util.rs b/src/util.rs
index fb0a56a..8c0f4be 100644
--- a/src/util.rs
+++ b/src/util.rs
@@ -376,6 +376,16 @@ impl Ark {
                     .otherwise(col("company"))
                     .alias("company"),
             ])
+            .with_columns(vec![
+                when(col("company").eq(lit("ARK BITCOIN ETF HOLDCO (ARKF)")))
+                    .then(lit("ARKB"))
+                    .otherwise(col("ticker"))
+                    .alias("ticker"),
+                when(col("company").eq(lit("ARK BITCOIN ETF HOLDCO (ARKF)")))
+                    .then(lit("ARKB"))
+                    .otherwise(col("company"))
+                    .alias("company"),
+            ])
             .collect()
         {
             df = x;
@@ -723,4 +733,40 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    #[serial]
+    fn arkf_format_arkb() -> Result<(), Error> {
+        let test_df = df![
+            "date" => ["2024-01-01", "2024-01-02"],
+            "ticker" => [None::<&str>, Some("TSLA")],
+            "cusip" => ["123abc", "TESLA"],
+            "company" => ["ARK BITCOIN ETF HOLDCO (ARKF)", "TESLA"],
+            "market_value" => [100, 400],
+            "shares" => [10, 20],
+            "share_price" => [10, 20],
+            "weight" => [10.00, 20.00]
+        ]?;
+
+        Ark::write_df_parquet("data/test/ARKF.parquet".into(), test_df.clone().into())?;
+        let read = Ark::new(Source::Read, Ticker::ARKF, Some("data/test".to_owned()))?.collect()?;
+        fs::remove_file("data/test/ARKF.parquet")?;
+
+        let df = Ark::df_format(read.into())?.collect()?;
+        assert_eq!(
+            df,
+            df![
+                "date" => ["2024-01-01", "2024-01-02"],
+                "ticker" => ["ARKB", "TSLA"],
+                "cusip" => ["123abc", "TESLA"],
+                "company" => ["ARKB", "TESLA"],
+                "market_value" => [100, 400],
+                "shares" => [10, 20],
+                "share_price" => [10, 20],
+                "weight" => [10.00, 20.00]
+            ]?
+        );
+
+        Ok(())
+    }
 }