1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
|
import unittest
from unittest.mock import patch, PropertyMock
from app import data_store
class TestDataStore(unittest.TestCase):
@patch.object(data_store.DataStore, "_monitor_objects")
@patch('app.data_store._write_to_gcs')
def test_save(self, mock_write, mock_monitor):
def add_ham():
# Adding two seen "ham" object to the internal store
ds._questionable_ham[0.1] = {"a": 1}
ds._questionable_ham[0.2] = {"b": 2}
ds = data_store.DataStore(spammable_type="foobar")
# save a spam object
ds.save(spammable={}, confidence=ds.SPAM_HAM_THRESHOLD + 0.1)
self.assertEqual(-1, ds._ham_spam_count, "The spam ham count was not decremented")
# Save another spam, this should save the two ham objects
add_ham()
ds.save(spammable={}, confidence=ds.SPAM_HAM_THRESHOLD + 0.1)
self.assertEqual(0, ds._ham_spam_count, "The spam ham count should be 0")
mock_write.assert_any_call(ds.gcs_file_path, {"a": 1})
mock_write.assert_any_call(ds.gcs_file_path, {"b": 2})
# Set the counts so there is an imbalance of ham in the labeled dataset
ds._ham_spam_count = -2
ds._labeled_count = 10
add_ham()
ds.save(spammable={}, confidence=ds.SPAM_HAM_THRESHOLD + 0.1)
self.assertEqual(-3, ds._ham_spam_count, "We should not have saved any ham")
# Set the counts so there is an imbalance of spam in the labeled dataset
ds._ham_spam_count = 1
ds._labeled_count = -5
add_ham()
ds.save(spammable={}, confidence=ds.SPAM_HAM_THRESHOLD + 0.1)
self.assertEqual(0, ds._ham_spam_count, "Score should have decremented by one when _ham_spam_count is positive and spam was saved")
self.assertEqual(-3, ds._labeled_count, "We should have saved two ham issues and incremented the _labeled_count")
|
