unit testing polars gc reenable

crates/nu_plugin_polars/src/cache/list.rs

@@ -16,7 +16,7 @@ impl PluginCommand for ListDF {
     }
 
     fn usage(&self) -> &str {
-        "Lists stored dataframes."
+        "Lists stored polars objects."
     }
 
     fn signature(&self) -> Signature {
@@ -119,6 +119,20 @@ impl PluginCommand for ListDF {
                     },
                     call.head,
                 ))),
+                PolarsPluginObject::NuPolarsTestData(_, _) => Ok(Some(Value::record(
+                    record! {
+                        "key" => Value::string(key.to_string(), call.head),
+                        "columns" => Value::nothing(call.head),
+                        "rows" => Value::nothing(call.head),
+                        "type" => Value::string("When", call.head),
+                        "estimated_size" => Value::nothing(call.head),
+                        "span_contents" =>  Value::string(span_contents.to_string(), call.head),
+                        "span_start" => Value::int(call.head.start as i64, call.head),
+                        "span_end" => Value::int(call.head.end as i64, call.head),
+                        "reference_count" => Value::int(value.reference_count as i64, call.head),
+                    },
+                    call.head,
+                )))
             }
         })?;
         let vals = vals.into_iter().flatten().collect();

crates/nu_plugin_polars/src/cache/mod.rs

@@ -172,3 +172,83 @@ pub(crate) fn cache_commands() -> Vec<Box<dyn PluginCommand<Plugin = PolarsPlugi
         Box::new(get::CacheGet),
     ]
 }
+
+#[cfg(test)]
+mod test {
+    use std::{cell::RefCell, rc::Rc};
+
+    use super::*;
+
+    struct MockEngineWrapper {
+        gc_enabled: Rc<RefCell<bool>>,
+    }
+
+    impl MockEngineWrapper {
+        fn new(gc_enabled: bool) -> Self {
+            Self {
+                gc_enabled: Rc::new(RefCell::new(gc_enabled)),
+            }
+        }
+
+        fn gc_enabled(&self) -> bool {
+            *self.gc_enabled.borrow()
+        }
+    }
+
+    impl EngineWrapper for &MockEngineWrapper {
+        fn get_env_var(&self, _key: &str) -> Option<String> {
+            todo!()
+        }
+
+        fn use_color(&self) -> bool {
+            todo!()
+        }
+
+        fn set_gc_disabled(&self, disabled: bool) -> Result<(), ShellError> {
+            let _ = self.gc_enabled.replace(!disabled);
+            Ok(())
+        }
+    }
+
+    #[test]
+    pub fn test_remove_plugin_cache_enable() {
+        let mock_engine = MockEngineWrapper::new(false);
+
+        let cache = Cache::default();
+        let mut lock = cache.cache.lock().expect("should be able to acquire lock");
+        let key0 = Uuid::new_v4();
+        lock.insert(
+            key0,
+            CacheValue {
+                uuid: Uuid::new_v4(),
+                value: PolarsPluginObject::NuPolarsTestData(Uuid::new_v4(), "object_0".into()),
+                created: Local::now().into(),
+                span: Span::unknown(),
+                reference_count: 1,
+            },
+        );
+
+        let key1 = Uuid::new_v4();
+        lock.insert(
+            key1,
+            CacheValue {
+                uuid: Uuid::new_v4(),
+                value: PolarsPluginObject::NuPolarsTestData(Uuid::new_v4(), "object_1".into()),
+                created: Local::now().into(),
+                span: Span::unknown(),
+                reference_count: 1,
+            },
+        );
+        drop(lock);
+
+        let _ = cache
+            .remove(&mock_engine, &key0, false)
+            .expect("should be able to remove key0");
+        assert!(!mock_engine.gc_enabled());
+
+        let _ = cache
+            .remove(&mock_engine, &key1, false)
+            .expect("should be able to remove key1");
+        assert!(mock_engine.gc_enabled());
+    }
+}

crates/nu_plugin_polars/src/dataframe/values/mod.rs

@@ -27,6 +27,7 @@ pub enum PolarsPluginType {
     NuExpression,
     NuLazyGroupBy,
     NuWhen,
+    NuPolarsTestData,
 }
 
 impl fmt::Display for PolarsPluginType {
@@ -37,6 +38,7 @@ impl fmt::Display for PolarsPluginType {
             Self::NuExpression => write!(f, "NuExpression"),
             Self::NuLazyGroupBy => write!(f, "NuLazyGroupBy"),
             Self::NuWhen => write!(f, "NuWhen"),
+            Self::NuPolarsTestData => write!(f, "NuPolarsTestData"),
         }
     }
 }
@@ -48,6 +50,7 @@ pub enum PolarsPluginObject {
     NuExpression(NuExpression),
     NuLazyGroupBy(NuLazyGroupBy),
     NuWhen(NuWhen),
+    NuPolarsTestData(Uuid, String),
 }
 
 impl PolarsPluginObject {
@@ -95,6 +98,7 @@ impl PolarsPluginObject {
             Self::NuExpression(_) => PolarsPluginType::NuExpression,
             Self::NuLazyGroupBy(_) => PolarsPluginType::NuLazyGroupBy,
             Self::NuWhen(_) => PolarsPluginType::NuWhen,
+            Self::NuPolarsTestData(_, _) => PolarsPluginType::NuPolarsTestData,
         }
     }
 
@@ -105,6 +109,7 @@ impl PolarsPluginObject {
             PolarsPluginObject::NuExpression(e) => e.id,
             PolarsPluginObject::NuLazyGroupBy(lg) => lg.id,
             PolarsPluginObject::NuWhen(w) => w.id,
+            PolarsPluginObject::NuPolarsTestData(id, _) => *id,
         }
     }
 
@@ -115,6 +120,9 @@ impl PolarsPluginObject {
             PolarsPluginObject::NuExpression(e) => e.into_value(span),
             PolarsPluginObject::NuLazyGroupBy(lg) => lg.into_value(span),
             PolarsPluginObject::NuWhen(w) => w.into_value(span),
+            PolarsPluginObject::NuPolarsTestData(id, s) => {
+                Value::string(format!("{id}:{s}"), Span::test_data())
+            }
         }
     }
 }