From 1da30bd0096d6fc2f70b70ec046873c1e5074d6a Mon Sep 17 00:00:00 2001
From: Pranav Anbarasu <pranavanba@gmail.com>
Date: Fri, 27 Oct 2023 00:15:54 +0000
Subject: [PATCH] Add unit tests for functions

---
 tests/testthat/test-filtering.R | 38 +++++++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/tests/testthat/test-filtering.R b/tests/testthat/test-filtering.R
index e69de29..f3c7cca 100644
--- a/tests/testthat/test-filtering.R
+++ b/tests/testthat/test-filtering.R
@@ -0,0 +1,38 @@
+library(testthat)
+
+test_that("dob2age correctly calculates age from Date of Birth", {
+  test_dataset <- data.frame(
+    date_of_birth = as.Date(c("1990-01-15", "1985-05-10", "2000-12-30"))
+  )
+  
+  arrow::write_dataset(test_dataset, path = "test_dob2age")
+  
+  dob2age("test_dob2age", "date_of_birth", input = ".", partitions = NULL)
+  
+  modified_dataset <- arrow::open_dataset("test_dob2age")
+  
+  expect_true("age" %in% names(modified_dataset))
+  expect_equal((modified_dataset %>% collect %>% pull(age)), c(33, 38, 23))
+  
+  unlink("test_dob2age/part-0.parquet")
+})
+
+test_that("drop_cols_datasets correctly drops specified columns", {
+  test_dataset <- data.frame(
+    column1 = c(1, 2, 3),
+    column2 = c("A", "B", "C"),
+    column3 = c(0.1, 0.2, 0.3)
+  )
+  
+  arrow::write_dataset(test_dataset, path = "test_drop_cols")
+  
+  columns_to_drop <- c("column1", "column2")
+  
+  drop_cols_datasets("test_drop_cols", columns = columns_to_drop, input = ".", output = ".")
+  
+  modified_dataset <- arrow::open_dataset("test_drop_cols")
+  
+  expect_true(all(!names(modified_dataset) %in% columns_to_drop))
+  
+  unlink("test_drop_cols/part-0.parquet")
+})