Refactoring Script functions; move cmd build to its own functions

Signed-off-by: Maxim Zhiburt <zhiburt@gmail.com>
redis-rs · May 11, 2022 · 50a48b3 · 50a48b3
1 parent fb861af
commit 50a48b3
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 48 deletions.
diff --git a/src/script.rs b/src/script.rs
@@ -4,6 +4,7 @@ use sha1_smol::Sha1;
 use crate::cmd::cmd;
 use crate::connection::ConnectionLike;
 use crate::types::{ErrorKind, FromRedisValue, RedisResult, ToRedisArgs};
+use crate::Cmd;
 
 /// Represents a lua script.
 #[derive(Debug, Clone)]
@@ -124,23 +125,12 @@ impl<'a> ScriptInvocation<'a> {
     /// Invokes the script and returns the result.
     #[inline]
     pub fn invoke<T: FromRedisValue>(&self, con: &mut dyn ConnectionLike) -> RedisResult<T> {
-        let mut eval_cmd = cmd("EVALSHA");
-        eval_cmd
-            .arg(self.script.hash.as_bytes())
-            .arg(self.keys.len())
-            .arg(&*self.keys)
-            .arg(&*self.args);
-
+        let eval_cmd = self.eval_cmd();
         match eval_cmd.query(con) {
             Ok(val) => Ok(val),
             Err(err) => {
                 if err.kind() == ErrorKind::NoScriptError {
-                    let mut load_cmd = cmd("SCRIPT");
-                    load_cmd
-                        .arg("LOAD")
-                        .arg(self.script.code.as_bytes())
-                        .query(con)?;
-
+                    self.load_cmd().query(con)?;
                     eval_cmd.query(con)
                 } else {
                     Err(err)
@@ -157,13 +147,7 @@ impl<'a> ScriptInvocation<'a> {
         C: crate::aio::ConnectionLike,
         T: FromRedisValue,
     {
-        let mut eval_cmd = cmd("EVALSHA");
-        eval_cmd
-            .arg(self.script.hash.as_bytes())
-            .arg(self.keys.len())
-            .arg(&*self.keys)
-            .arg(&*self.args);
-
+        let eval_cmd = self.eval_cmd();
         match eval_cmd.query_async(con).await {
             Ok(val) => {
                 // Return the value from the script evaluation
@@ -172,10 +156,7 @@ impl<'a> ScriptInvocation<'a> {
             Err(err) => {
                 // Load the script into Redis if the script hash wasn't there already
                 if err.kind() == ErrorKind::NoScriptError {
-                    let mut load_cmd = cmd("SCRIPT");
-                    load_cmd.arg("LOAD").arg(self.script.code.as_bytes());
-
-                    load_cmd.query_async(con).await?;
+                    self.load_cmd().query_async(con).await?;
                     eval_cmd.query_async(con).await
                 } else {
                     Err(err)
@@ -187,10 +168,7 @@ impl<'a> ScriptInvocation<'a> {
     /// Loads the script and returns the SHA1 of it.
     #[inline]
     pub fn load(&self, con: &mut dyn ConnectionLike) -> RedisResult<String> {
-        let hash: String = cmd("SCRIPT")
-            .arg("LOAD")
-            .arg(self.script.code.as_bytes())
-            .query(con)?;
+        let hash: String = self.load_cmd().query(con)?;
 
         debug_assert_eq!(hash, self.script.hash);
 
@@ -204,13 +182,25 @@ impl<'a> ScriptInvocation<'a> {
     where
         C: crate::aio::ConnectionLike,
     {
-        let hash: String = cmd("SCRIPT")
-            .arg("LOAD")
-            .arg(self.script.code.as_bytes())
-            .query_async(con).await?;
+        let hash: String = self.load_cmd().query_async(con).await?;
 
         debug_assert_eq!(hash, self.script.hash);
 
         Ok(hash)
     }
+
+    fn load_cmd(&self) -> Cmd {
+        let mut cmd = cmd("SCRIPT");
+        cmd.arg("LOAD").arg(self.script.code.as_bytes());
+        cmd
+    }
+
+    fn eval_cmd(&self) -> Cmd {
+        let mut cmd = cmd("EVALSHA");
+        cmd.arg(self.script.hash.as_bytes())
+            .arg(self.keys.len())
+            .arg(&*self.keys)
+            .arg(&*self.args);
+        cmd
+    }
 }
diff --git a/tests/test_async.rs b/tests/test_async.rs
@@ -354,19 +354,14 @@ fn test_script() {
 #[cfg(feature = "script")]
 fn test_script_load() {
     let ctx = TestContext::new();
-    let mut con = ctx.connection();
-
     let script = redis::Script::new("return 'Hello World'");
 
     block_on_all(async move {
-        let mut con = ctx.multiplexed_async_connection().await?;
+        let mut con = ctx.multiplexed_async_connection().await.unwrap();
 
-        let hash = script.prepare_invoke().load_async(&mut con);
-        assert_eq!(hash, Ok(script.get_hash().to_string()));
-
-        Ok(())
-    })
-    .unwrap();
+        let hash = script.prepare_invoke().load_async(&mut con).await.unwrap();
+        assert_eq!(hash, script.get_hash().to_string());
+    });
 }
 
 #[test]

diff --git a/tests/test_async_async_std.rs b/tests/test_async_async_std.rs
@@ -295,14 +295,11 @@ fn test_script_load() {
     let script = redis::Script::new("return 'Hello World'");
 
     block_on_all(async move {
-        let mut con = ctx.multiplexed_async_connection_async_std().await?;
-
-        let hash = script.prepare_invoke().load_async(&mut con);
-        assert_eq!(hash, Ok(script.get_hash().to_string()));
+        let mut con = ctx.multiplexed_async_connection_async_std().await.unwrap();
 
-        Ok(())
-    })
-    .unwrap();
+        let hash = script.prepare_invoke().load_async(&mut con).await.unwrap();
+        assert_eq!(hash, script.get_hash().to_string());
+    });
 }
 
 #[test]