From ed9d148411744ea543dccaec9247576cbdedf90d Mon Sep 17 00:00:00 2001
From: simon987 <me@simon987.net>
Date: Thu, 21 Oct 2021 19:52:42 -0400
Subject: [PATCH] Fix tests

---
 .gitignore                 |  5 ++++-
 hexlib/db.py               | 14 ++++++++++----
 setup.py                   |  3 ++-
 test/__init__.py           |  0
 test/test_VolatileState.py |  9 +++++++++
 test/test_download_file.py | 10 +++++++---
 6 files changed, 32 insertions(+), 9 deletions(-)
 create mode 100644 test/__init__.py

diff --git a/.gitignore b/.gitignore
index 431af02..bb63df1 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,4 +1,7 @@
 *.iml
 .idea/
 *.db
-*.png
\ No newline at end of file
+*.png
+hexlib.egg-info
+build/
+dist/
\ No newline at end of file
diff --git a/hexlib/db.py b/hexlib/db.py
index abe111f..1fac624 100644
--- a/hexlib/db.py
+++ b/hexlib/db.py
@@ -36,6 +36,9 @@ class VolatileState:
     def __getitem__(self, table):
         return RedisTable(self, table, self._sep)
 
+    def __delitem__(self, key):
+        self.rdb.delete(f"{self.prefix}{self._sep}{key}")
+
 
 class VolatileQueue:
     """Quick and dirty volatile queue-like redis wrapper"""
@@ -68,6 +71,9 @@ class VolatileBooleanState:
     def __getitem__(self, table):
         return RedisBooleanTable(self, table, self._sep)
 
+    def __delitem__(self, table):
+        self.rdb.delete(f"{self.prefix}{self._sep}{table}")
+
 
 class RedisTable:
     def __init__(self, state, table, sep=""):
@@ -89,9 +95,9 @@ class RedisTable:
         self._state.rdb.hdel(self._key, str(key))
 
     def __iter__(self):
-        val = self._state.rdb.hgetall(self._key)
-        if val:
-            return ((k, umsgpack.loads(v)) for k, v in val.items())
+        for val in self._state.rdb.hscan(self._key):
+            if val:
+                return ((k, umsgpack.loads(v)) for k, v in val.items())
 
 
 class RedisBooleanTable:
@@ -114,7 +120,7 @@ class RedisBooleanTable:
         self._state.rdb.srem(self._key, str(key))
 
     def __iter__(self):
-        return iter(self._state.rdb.smembers(self._key))
+        yield from self._state.rdb.sscan_iter(self._key)
 
 
 class Table:
diff --git a/setup.py b/setup.py
index 39139c3..7da3eb3 100644
--- a/setup.py
+++ b/setup.py
@@ -14,6 +14,7 @@ setup(
     install_requires=[
         "ImageHash", "influxdb", "siphash", "python-dateutil", "redis", "orjson", "zstandard",
         "u-msgpack-python", "psycopg2-binary", "bs4", "lxml", "nltk", "numpy",
-        "matplotlib", "scikit-learn", "fake-useragent @ git+git://github.com/Jordan9675/fake-useragent"
+        "matplotlib", "scikit-learn", "fake-useragent @ git+git://github.com/Jordan9675/fake-useragent",
+        "requests"
     ]
 )
diff --git a/test/__init__.py b/test/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/test/test_VolatileState.py b/test/test_VolatileState.py
index cb0933b..c0c4185 100644
--- a/test/test_VolatileState.py
+++ b/test/test_VolatileState.py
@@ -1,10 +1,15 @@
 from unittest import TestCase
 
 from hexlib.db import VolatileState, VolatileBooleanState, VolatileQueue
+from hexlib.env import get_redis
 
 
 class TestVolatileState(TestCase):
 
+    def setUp(self) -> None:
+        rdb = get_redis()
+        rdb.delete("test1a", "test1b", "test1c", "test1:a", "test2b")
+
     def test_get_set(self):
         s = VolatileState(prefix="test1")
         val = {
@@ -53,6 +58,10 @@ class TestVolatileState(TestCase):
 
 class TestVolatileBoolState(TestCase):
 
+    def setUp(self) -> None:
+        rdb = get_redis()
+        rdb.delete("test1a", "test1b", "test1c", "test1:a", "test2b")
+
     def test_get_set(self):
         s = VolatileBooleanState(prefix="test1")
 
diff --git a/test/test_download_file.py b/test/test_download_file.py
index 28f4fc7..61dee11 100644
--- a/test/test_download_file.py
+++ b/test/test_download_file.py
@@ -2,12 +2,16 @@ import os
 from unittest import TestCase
 
 from hexlib.web import download_file
+import warnings
 
 
 class TestDownloadFile(TestCase):
 
+    def setUp(self) -> None:
+        warnings.filterwarnings(action="ignore", category=ResourceWarning)
+
     def test_download_file(self):
-        download_file("http://ovh.net/files/10Mb.dat", "/tmp/10Mb.dat")
+        download_file("https://github.com/simon987/hexlib/raw/master/10MB.bin", "/tmp/10Mb.dat")
         self.assertTrue(os.path.exists("/tmp/10Mb.dat"))
         os.remove("/tmp/10Mb.dat")
 
@@ -22,8 +26,8 @@ class TestDownloadFile(TestCase):
         self.assertEqual(len(exceptions), 3)
 
     def test_download_file_meta(self):
-        download_file("http://ovh.net/files/10Mb.dat", "/tmp/10Mb.dat", save_meta=True)
+        download_file("https://github.com/simon987/hexlib/raw/master/10MB.bin", "/tmp/10Mb.dat", save_meta=True)
         self.assertTrue(os.path.exists("/tmp/10Mb.dat"))
         self.assertTrue(os.path.exists("/tmp/10Mb.dat.meta"))
         os.remove("/tmp/10Mb.dat")
-        # os.remove("/tmp/10Mb.dat.meta")
+        os.remove("/tmp/10Mb.dat.meta")