KaneAI - World’s First E2E Software Testing Agent.Book a Demo 
How to use mock_client method in tempest
Best Python code snippet using tempest_python
magics_impl_test.py
Source:magics_impl_test.py 
1#!/usr/bin/env python2import collections3import io4import ipaddress5import time6from unittest import mock7from absl.testing import absltest8import pandas as pd9import grr_colab10from grr_colab import _timeout11from grr_colab import magics_impl12from grr_response_proto import artifact_pb213from grr_response_proto.api import client_pb214class GrrSetNoFlowTimeoutImplTest(absltest.TestCase):15 def testNone(self):16 with mock.patch.object(_timeout, '_FLOW_TIMEOUT', 30):17 magics_impl.grr_set_no_flow_timeout_impl()18 self.assertIsNone(_timeout._FLOW_TIMEOUT)19class GrrSetDefaultFlowTimeoutImplTest(absltest.TestCase):20 def testDefault(self):21 with mock.patch.object(_timeout, '_FLOW_TIMEOUT', 10):22 magics_impl.grr_set_default_flow_timeout_impl()23 self.assertEqual(_timeout._FLOW_TIMEOUT, 30)24class GrrSetFlowTimeoutImplTest(absltest.TestCase):25 def testNoneValueError(self):26 with mock.patch.object(_timeout, '_FLOW_TIMEOUT', 30):27 with self.assertRaises(ValueError):28 magics_impl.grr_set_flow_timeout_impl(None)29 def testZero(self):30 with mock.patch.object(_timeout, '_FLOW_TIMEOUT', 30):31 magics_impl.grr_set_flow_timeout_impl(0)32 self.assertEqual(_timeout._FLOW_TIMEOUT, 0)33 def testNonZero(self):34 with mock.patch.object(_timeout, '_FLOW_TIMEOUT', 30):35 magics_impl.grr_set_flow_timeout_impl(10)36 self.assertEqual(_timeout._FLOW_TIMEOUT, 10)37 def testNegativeValueError(self):38 with mock.patch.object(_timeout, '_FLOW_TIMEOUT', 30):39 with self.assertRaises(ValueError):40 magics_impl.grr_set_flow_timeout_impl(-10)41class GrrListArtifactsImplTest(absltest.TestCase):42 def testEmptyResults(self):43 with mock.patch.object(44 grr_colab, 'list_artifacts', return_value=[]) as mock_fn:45 magics_impl.grr_list_artifacts_impl()46 mock_fn.assert_called_once_with()47 def testPriorityColumns(self):48 artifact = artifact_pb2.ArtifactDescriptor()49 artifact.artifact.name = 'foo'50 artifact.artifact.doc = 'bar'51 artifact.is_custom = True52 with mock.patch.object(53 grr_colab, 'list_artifacts', return_value=[artifact]):54 df = magics_impl.grr_list_artifacts_impl()55 self.assertEqual((1, 3), df.shape)56 self.assertEqual(57 list(df.columns), ['artifact.name', 'artifact.doc', 'is_custom'])58 self.assertEqual(df['artifact.name'][0], 'foo')59 self.assertEqual(df['artifact.doc'][0], 'bar')60 self.assertTrue(df['is_custom'][0])61class GrrSearchClientsImplTest(absltest.TestCase):62 def testNoArgs(self):63 with mock.patch.object(64 grr_colab.Client, 'search', return_value=[]) as mock_fn:65 magics_impl.grr_search_clients_impl()66 mock_fn.assert_called_once_with(67 version=None, host=None, labels=None, mac=None, ip=None, user=None)68 def testWithArgs(self):69 with mock.patch.object(70 grr_colab.Client, 'search', return_value=[]) as mock_fn:71 magics_impl.grr_search_clients_impl(72 version='test_v',73 host='test_host',74 labels=['test_label'],75 mac='test_mac',76 ip='test_ip',77 user='test_user')78 mock_fn.assert_called_once_with(79 version='test_v',80 host='test_host',81 labels=['test_label'],82 mac='test_mac',83 ip='test_ip',84 user='test_user')85 def testIsDataframe(self):86 mock_client = _MockClient()87 with mock.patch.object(88 grr_colab.Client, 'search', return_value=[mock_client]):89 df = magics_impl.grr_search_clients_impl()90 self.assertIsInstance(df, pd.DataFrame)91 self.assertEqual(df.shape[0], 1)92 def testSortedByLastSeen(self):93 mock_clients = [94 _MockClient(client_id='foo', last_seen_at=1000),95 _MockClient(client_id='bar', last_seen_at=10),96 _MockClient(client_id='quux', last_seen_at=100),97 ]98 with mock.patch.object(99 grr_colab.Client, 'search', return_value=mock_clients):100 df = magics_impl.grr_search_clients_impl()101 self.assertEqual(df.shape[0], 3)102 self.assertEqual(list(df['client_id']), ['foo', 'quux', 'bar'])103 def testOnlineColumns(self):104 current_time_secs = int(time.time())105 mock_clients = [106 _MockClient(last_seen_at=(current_time_secs - 2 * 60) * (10**6)),107 _MockClient(last_seen_at=(current_time_secs - 3 * 60 * 60) * (10**6)),108 _MockClient(109 last_seen_at=(current_time_secs - 4 * 60 * 60 * 24) * (10**6)),110 ]111 with mock.patch.object(112 grr_colab.Client, 'search', return_value=mock_clients):113 df = magics_impl.grr_search_clients_impl()114 self.assertEqual(df.shape[0], 3)115 self.assertIn('online', df.columns)116 self.assertEqual(df['online'][0], 'online')117 self.assertEqual(df['online'][1], 'seen-1d')118 self.assertEqual(df['online'][2], 'offline')119 self.assertIn('online.pretty', df.columns)120 self.assertEqual(df['online.pretty'][0], 'ð')121 self.assertEqual(df['online.pretty'][1], 'ð')122 self.assertEqual(df['online.pretty'][2], 'ð')123 def testLastSeenAgoColumn(self):124 current_time_secs = 1560000000125 mock_clients = [126 _MockClient(last_seen_at=(current_time_secs + 1) * (10**6)),127 _MockClient(last_seen_at=(current_time_secs - 1) * (10**6)),128 _MockClient(last_seen_at=(current_time_secs - 2 * 60) * (10**6)),129 _MockClient(last_seen_at=(current_time_secs - 3 * 60 * 60) * (10**6)),130 _MockClient(131 last_seen_at=(current_time_secs - 4 * 60 * 60 * 24) * (10**6)),132 ]133 with mock.patch.object(time, 'time', return_value=current_time_secs):134 with mock.patch.object(135 grr_colab.Client, 'search', return_value=mock_clients):136 df = magics_impl.grr_search_clients_impl()137 self.assertEqual(df.shape[0], 5)138 self.assertIn('last_seen_ago', df.columns)139 self.assertEqual(df['last_seen_ago'][0], 'in 1 seconds')140 self.assertEqual(df['last_seen_ago'][1], '1 seconds ago')141 self.assertEqual(df['last_seen_ago'][2], '2 minutes ago')142 self.assertEqual(df['last_seen_ago'][3], '3 hours ago')143 self.assertEqual(df['last_seen_ago'][4], '4 days ago')144 def testPriorityColumns(self):145 mock_client = _MockClient(client_id='foo')146 mock_client.set_hostname('test_hostname')147 mock_client.set_os_version('test_version')148 with mock.patch.object(149 grr_colab.Client, 'search', return_value=[mock_client]):150 df = magics_impl.grr_search_clients_impl()151 self.assertGreaterEqual(df.shape[1], 7)152 self.assertEqual(153 list(df)[:7], [154 'online.pretty', 'online', 'client_id', 'last_seen_ago',155 'last_seen_at.pretty', 'knowledge_base.fqdn', 'os_info.version'156 ])157class GrrSearchOnlineClientsImplTest(absltest.TestCase):158 def testContainsOnlineOnly(self):159 current_time_secs = int(time.time())160 mock_clients = [161 _MockClient(162 client_id='foo',163 last_seen_at=(current_time_secs - 2 * 60) * (10**6)),164 _MockClient(165 client_id='bar',166 last_seen_at=(current_time_secs - 3 * 60 * 60) * (10**6)),167 _MockClient(168 client_id='quux',169 last_seen_at=(current_time_secs - 4 * 60 * 60 * 24) * (10**6)),170 ]171 with mock.patch.object(172 grr_colab.Client, 'search', return_value=mock_clients):173 df = magics_impl.grr_search_online_clients_impl()174 self.assertEqual(df.shape[0], 1)175 self.assertEqual(df['client_id'][0], 'foo')176 self.assertEqual(df['online'][0], 'online')177 self.assertEqual(df['online.pretty'][0], 'ð')178 def testEmptyResults(self):179 current_time_secs = int(time.time())180 mock_clients = [181 _MockClient(last_seen_at=(current_time_secs - 3 * 60 * 60) * (10**6)),182 _MockClient(183 last_seen_at=(current_time_secs - 4 * 60 * 60 * 24) * (10**6)),184 ]185 with mock.patch.object(186 grr_colab.Client, 'search', return_value=mock_clients):187 df = magics_impl.grr_search_online_clients_impl()188 self.assertEqual(df.shape[0], 0)189 def testSortedByLastSeen(self):190 current_time_secs = int(time.time())191 mock_clients = [192 _MockClient(193 client_id='foo', last_seen_at=(current_time_secs - 5) * (10**6)),194 _MockClient(195 client_id='bar', last_seen_at=(current_time_secs - 3) * (10**6)),196 _MockClient(197 client_id='quux', last_seen_at=(current_time_secs - 10) * (10**6)),198 ]199 with mock.patch.object(200 grr_colab.Client, 'search', return_value=mock_clients):201 df = magics_impl.grr_search_online_clients_impl()202 self.assertEqual(df.shape[0], 3)203 self.assertEqual(list(df['client_id']), ['bar', 'foo', 'quux'])204class GrrSetClientImplTest(absltest.TestCase):205 def testBothHostnameAndId(self):206 with mock.patch.object(grr_colab.Client, 'with_id') as with_id_fn:207 with mock.patch.object(grr_colab.Client,208 'with_hostname') as with_hostname_fn:209 with mock.patch.object(magics_impl, '_state', magics_impl._State()):210 with self.assertRaises(ValueError):211 magics_impl.grr_set_client_impl('test_hostname', 'test_id')212 self.assertFalse(with_id_fn.called)213 self.assertFalse(with_hostname_fn.called)214 def testNeitherHostnameOrId(self):215 with mock.patch.object(grr_colab.Client, 'with_id') as with_id_fn:216 with mock.patch.object(grr_colab.Client,217 'with_hostname') as with_hostname_fn:218 with mock.patch.object(magics_impl, '_state', magics_impl._State()):219 with self.assertRaises(ValueError):220 magics_impl.grr_set_client_impl()221 self.assertFalse(with_id_fn.called)222 self.assertFalse(with_hostname_fn.called)223 def testWithHostname(self):224 mock_client = _MockClient(client_id='bar')225 with mock.patch.object(grr_colab.Client, 'with_id') as with_id_fn:226 with mock.patch.object(227 grr_colab.Client, 'with_hostname',228 return_value=mock_client) as with_hostname_fn:229 with mock.patch.object(magics_impl, '_state', magics_impl._State()):230 magics_impl.grr_set_client_impl(hostname='foo')231 self.assertFalse(with_id_fn.called)232 with_hostname_fn.assert_called_once_with('foo')233 self.assertEqual(magics_impl._state.client.id, mock_client.id)234 def testWithId(self):235 mock_client = _MockClient(client_id='foo')236 with mock.patch.object(237 grr_colab.Client, 'with_id', return_value=mock_client) as with_id_fn:238 with mock.patch.object(grr_colab.Client,239 'with_hostname') as with_hostname_fn:240 with mock.patch.object(magics_impl, '_state', magics_impl._State()):241 magics_impl.grr_set_client_impl(client='foo')242 self.assertFalse(with_hostname_fn.called)243 with_id_fn.assert_called_once_with('foo')244 self.assertEqual(magics_impl._state.client.id, mock_client.id)245class GrrRequestApprovalImplTest(absltest.TestCase):246 def testNoClientSelected(self):247 with self.assertRaises(magics_impl.NoClientSelectedError):248 magics_impl.grr_request_approval_impl('foo', ['bar'])249 def testNoWait(self):250 mock_client = mock.MagicMock()251 with mock.patch.object(magics_impl._state, 'client', mock_client):252 magics_impl.grr_request_approval_impl('foo', ['bar'])253 mock_client.request_approval.assert_called_once_with(254 reason='foo', approvers=['bar'])255 self.assertFalse(mock_client.request_approval_and_wait.called)256 def testWait(self):257 mock_client = mock.MagicMock()258 with mock.patch.object(magics_impl._state, 'client', mock_client):259 magics_impl.grr_request_approval_impl('foo', ['bar'], wait=True)260 mock_client.request_approval_and_wait.assert_called_once_with(261 reason='foo', approvers=['bar'])262 self.assertFalse(mock_client.request_approval.called)263class GrrIdImplTest(absltest.TestCase):264 def testNoClientSelected(self):265 with self.assertRaises(magics_impl.NoClientSelectedError):266 magics_impl.grr_id_impl()267 def testWithClientSet(self):268 mock_client = _MockClient(client_id='foo')269 with mock.patch.object(magics_impl._state, 'client', mock_client):270 client_id = magics_impl.grr_id_impl()271 self.assertEqual(client_id, 'foo')272class GrrCdImplTest(absltest.TestCase):273 def testNoClientSelected(self):274 with self.assertRaises(magics_impl.NoClientSelectedError):275 magics_impl.grr_cd_impl('foo')276 def testRelativePath(self):277 mock_client = _MockClient()278 with mock.patch.object(magics_impl._state, 'client', mock_client):279 with mock.patch.object(magics_impl._state, 'cur_dir', '/foo'):280 magics_impl.grr_cd_impl('bar/quux')281 self.assertEqual(magics_impl._state.cur_dir, '/foo/bar/quux')282 def testGoBack(self):283 mock_client = _MockClient()284 with mock.patch.object(magics_impl._state, 'client', mock_client):285 with mock.patch.object(magics_impl._state, 'cur_dir', '/foo/bar/quux'):286 magics_impl.grr_cd_impl('../..')287 self.assertEqual(magics_impl._state.cur_dir, '/foo')288 def testAbsolutePath(self):289 mock_client = _MockClient()290 with mock.patch.object(magics_impl._state, 'client', mock_client):291 with mock.patch.object(magics_impl._state, 'cur_dir', '/foo'):292 magics_impl.grr_cd_impl('/bar/quux')293 self.assertEqual(magics_impl._state.cur_dir, '/bar/quux')294class GrrPwdImplTest(absltest.TestCase):295 def testNoClientSelected(self):296 with self.assertRaises(magics_impl.NoClientSelectedError):297 magics_impl.grr_pwd_impl()298 def testRoot(self):299 mock_client = _MockClient()300 with mock.patch.object(magics_impl._state, 'client', mock_client):301 cur_dir = magics_impl.grr_pwd_impl()302 self.assertEqual(cur_dir, '/')303 def testNonRootPath(self):304 mock_client = _MockClient()305 with mock.patch.object(magics_impl._state, 'client', mock_client):306 with mock.patch.object(magics_impl._state, 'cur_dir', '/foo/bar'):307 cur_dir = magics_impl.grr_pwd_impl()308 self.assertEqual(cur_dir, '/foo/bar')309class GrrLsImplTest(absltest.TestCase):310 def testNoClientSelected(self):311 with self.assertRaises(magics_impl.NoClientSelectedError):312 magics_impl.grr_ls_impl('/foo')313 def testRelativePath(self):314 mock_client = mock.MagicMock()315 with mock.patch.object(magics_impl._state, 'client', mock_client):316 with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):317 magics_impl.grr_ls_impl('foo/bar')318 mock_client.os.ls.assert_called_once_with('/quux/foo/bar')319 self.assertFalse(mock_client.os.cached.ls.called)320 def testAbsolutePath(self):321 mock_client = mock.MagicMock()322 with mock.patch.object(magics_impl._state, 'client', mock_client):323 with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):324 magics_impl.grr_ls_impl('/foo/bar')325 mock_client.os.ls.assert_called_once_with('/foo/bar')326 self.assertFalse(mock_client.os.cached.ls.called)327 def testCurrentPath(self):328 mock_client = mock.MagicMock()329 with mock.patch.object(magics_impl._state, 'client', mock_client):330 with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):331 magics_impl.grr_ls_impl()332 mock_client.os.ls.assert_called_once_with('/quux')333 self.assertFalse(mock_client.os.cached.ls.called)334 def testCached(self):335 mock_client = mock.MagicMock()336 with mock.patch.object(magics_impl._state, 'client', mock_client):337 with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):338 magics_impl.grr_ls_impl('foo/bar', cached=True)339 mock_client.os.cached.ls.assert_called_once_with('/quux/foo/bar')340 self.assertFalse(mock_client.os.ls.called)341 def testTskPathType(self):342 mock_client = mock.MagicMock()343 with mock.patch.object(magics_impl._state, 'client', mock_client):344 with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):345 magics_impl.grr_ls_impl('foo/bar', path_type=magics_impl.TSK)346 mock_client.tsk.ls.assert_called_once_with('/quux/foo/bar')347 def testNtfsPathType(self):348 mock_client = mock.MagicMock()349 with mock.patch.object(magics_impl._state, 'client', mock_client):350 with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):351 magics_impl.grr_ls_impl('foo/bar', path_type=magics_impl.NTFS)352 mock_client.ntfs.ls.assert_called_once_with('/quux/foo/bar')353 def testRegistryPathType(self):354 mock_client = mock.MagicMock()355 with mock.patch.object(magics_impl._state, 'client', mock_client):356 with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):357 magics_impl.grr_ls_impl('foo/bar', path_type=magics_impl.REGISTRY)358 mock_client.registry.ls.assert_called_once_with('/quux/foo/bar')359class GrrStatImplTest(absltest.TestCase):360 def testNoClientSelected(self):361 with self.assertRaises(magics_impl.NoClientSelectedError):362 magics_impl.grr_stat_impl('foo')363 def testRelativePath(self):364 mock_client = mock.MagicMock()365 with mock.patch.object(magics_impl._state, 'client', mock_client):366 with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):367 magics_impl.grr_stat_impl('foo/bar')368 mock_client.os.glob.assert_called_once_with('/quux/foo/bar')369 def testAbsolutePath(self):370 mock_client = mock.MagicMock()371 with mock.patch.object(magics_impl._state, 'client', mock_client):372 with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):373 magics_impl.grr_stat_impl('/foo/bar')374 mock_client.os.glob.assert_called_once_with('/foo/bar')375 def testTskPathType(self):376 mock_client = mock.MagicMock()377 with mock.patch.object(magics_impl._state, 'client', mock_client):378 magics_impl.grr_stat_impl('/foo/bar', path_type=magics_impl.TSK)379 mock_client.tsk.glob.assert_called_once_with('/foo/bar')380 def testNtfsPathType(self):381 mock_client = mock.MagicMock()382 with mock.patch.object(magics_impl._state, 'client', mock_client):383 magics_impl.grr_stat_impl('/foo/bar', path_type=magics_impl.NTFS)384 mock_client.ntfs.glob.assert_called_once_with('/foo/bar')385 def testRegistryPathType(self):386 mock_client = mock.MagicMock()387 with mock.patch.object(magics_impl._state, 'client', mock_client):388 magics_impl.grr_stat_impl('/foo/bar', path_type=magics_impl.REGISTRY)389 mock_client.registry.glob.assert_called_once_with('/foo/bar')390class GrrHeadImplTest(absltest.TestCase):391 def testNoClientSelected(self):392 with self.assertRaises(magics_impl.NoClientSelectedError):393 magics_impl.grr_head_impl('foo')394 def testRelativePath(self):395 mock_client = _MockClient()396 mock_client.os.add_file('/quux/foo/bar', b'foo bar')397 mock_client.os.cached.add_file('/quux/foo/bar', b'bar foo')398 with mock.patch.object(magics_impl._state, 'client', mock_client):399 with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):400 data = magics_impl.grr_head_impl('foo/bar')401 self.assertEqual(data, b'foo bar')402 def testAbsolutePath(self):403 mock_client = _MockClient()404 mock_client.os.add_file('/foo/bar', b'foo bar')405 mock_client.os.cached.add_file('/foo/bar', b'bar foo')406 with mock.patch.object(magics_impl._state, 'client', mock_client):407 with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):408 data = magics_impl.grr_head_impl('/foo/bar')409 self.assertEqual(data, b'foo bar')410 def testCached(self):411 mock_client = _MockClient()412 mock_client.os.add_file('/foo/bar', b'foo bar')413 mock_client.os.cached.add_file('/foo/bar', b'bar foo')414 with mock.patch.object(magics_impl._state, 'client', mock_client):415 with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):416 data = magics_impl.grr_head_impl('/foo/bar', cached=True)417 self.assertEqual(data, b'bar foo')418 def testTskPathType(self):419 mock_client = _MockClient()420 mock_client.tsk.add_file('/foo/bar', b'foo bar')421 with mock.patch.object(magics_impl._state, 'client', mock_client):422 data = magics_impl.grr_head_impl('/foo/bar', path_type=magics_impl.TSK)423 self.assertEqual(data, b'foo bar')424 def testNtfsPathType(self):425 mock_client = _MockClient()426 mock_client.ntfs.add_file('/foo/bar', b'foo bar')427 with mock.patch.object(magics_impl._state, 'client', mock_client):428 data = magics_impl.grr_head_impl('/foo/bar', path_type=magics_impl.NTFS)429 self.assertEqual(data, b'foo bar')430 def testRegistryPathType(self):431 mock_client = _MockClient()432 mock_client.registry.add_file('/foo/bar', b'foo bar')433 with mock.patch.object(magics_impl._state, 'client', mock_client):434 data = magics_impl.grr_head_impl(435 '/foo/bar', path_type=magics_impl.REGISTRY)436 self.assertEqual(data, b'foo bar')437 def testBytes(self):438 mock_client = _MockClient()439 mock_client.os.add_file('/foo/bar', b'foo bar')440 mock_client.os.cached.add_file('/foo/bar', b'bar foo')441 with mock.patch.object(magics_impl._state, 'client', mock_client):442 with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):443 data = magics_impl.grr_head_impl('/foo/bar', bytes=3)444 self.assertEqual(data, b'foo')445 def testOffset(self):446 mock_client = _MockClient()447 mock_client.os.add_file('/foo/bar', b'foo bar')448 mock_client.os.cached.add_file('/foo/bar', b'bar foo')449 with mock.patch.object(magics_impl._state, 'client', mock_client):450 with mock.patch.object(magics_impl._state, 'cur_dir', '/'):451 data = magics_impl.grr_head_impl('/foo/bar', bytes=3, offset=4)452 self.assertEqual(data, b'bar')453class GrrGrepImplTest(absltest.TestCase):454 def testNoClientSelected(self):455 with self.assertRaises(magics_impl.NoClientSelectedError):456 magics_impl.grr_grep_impl('foo', 'bar')457 def testRelativePath(self):458 mock_client = mock.MagicMock()459 with mock.patch.object(magics_impl._state, 'client', mock_client):460 with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):461 magics_impl.grr_grep_impl('foo bar', 'foo/bar')462 mock_client.os.grep.assert_called_once_with('/quux/foo/bar', b'foo bar')463 def testAbsolutePath(self):464 mock_client = mock.MagicMock()465 with mock.patch.object(magics_impl._state, 'client', mock_client):466 with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):467 magics_impl.grr_grep_impl('foo bar', '/foo/bar')468 mock_client.os.grep.assert_called_once_with('/foo/bar', b'foo bar')469 def testFixedStrings(self):470 mock_client = mock.MagicMock()471 with mock.patch.object(magics_impl._state, 'client', mock_client):472 with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):473 magics_impl.grr_grep_impl('foo bar', '/foo/bar', fixed_strings=True)474 mock_client.os.fgrep.assert_called_once_with('/foo/bar', b'foo bar')475 self.assertFalse(mock_client.grep.called)476 def testTskPathType(self):477 mock_client = mock.MagicMock()478 with mock.patch.object(magics_impl._state, 'client', mock_client):479 magics_impl.grr_grep_impl(480 'foo bar', '/foo/bar', path_type=magics_impl.TSK)481 mock_client.tsk.grep.assert_called_once_with('/foo/bar', b'foo bar')482 def testNtfsPathType(self):483 mock_client = mock.MagicMock()484 with mock.patch.object(magics_impl._state, 'client', mock_client):485 magics_impl.grr_grep_impl(486 'foo bar', '/foo/bar', path_type=magics_impl.NTFS)487 mock_client.ntfs.grep.assert_called_once_with('/foo/bar', b'foo bar')488 def testRegistryPathType(self):489 mock_client = mock.MagicMock()490 with mock.patch.object(magics_impl._state, 'client', mock_client):491 magics_impl.grr_grep_impl(492 'foo bar', '/foo/bar', path_type=magics_impl.REGISTRY)493 mock_client.registry.grep.assert_called_once_with('/foo/bar', b'foo bar')494 def testHexString(self):495 mock_client = mock.MagicMock()496 with mock.patch.object(magics_impl._state, 'client', mock_client):497 magics_impl.grr_grep_impl('ffac90', '/foo/bar', hex_string=True)498 mock_client.os.grep.assert_called_once_with('/foo/bar', b'\xff\xac\x90')499class GrrFgrepImplTest(absltest.TestCase):500 def testNoClientSelected(self):501 with self.assertRaises(magics_impl.NoClientSelectedError):502 magics_impl.grr_fgrep_impl('foo', 'bar')503 def testRelativePath(self):504 mock_client = mock.MagicMock()505 with mock.patch.object(magics_impl._state, 'client', mock_client):506 with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):507 magics_impl.grr_fgrep_impl('foo bar', 'foo/bar')508 mock_client.os.fgrep.assert_called_once_with('/quux/foo/bar',509 b'foo bar')510 def testAbsolutePath(self):511 mock_client = mock.MagicMock()512 with mock.patch.object(magics_impl._state, 'client', mock_client):513 with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):514 magics_impl.grr_fgrep_impl('foo bar', '/foo/bar')515 mock_client.os.fgrep.assert_called_once_with('/foo/bar', b'foo bar')516 def testTskPathType(self):517 mock_client = mock.MagicMock()518 with mock.patch.object(magics_impl._state, 'client', mock_client):519 magics_impl.grr_fgrep_impl(520 'foo bar', '/foo/bar', path_type=magics_impl.TSK)521 mock_client.tsk.fgrep.assert_called_once_with('/foo/bar', b'foo bar')522 def testNtfsPathType(self):523 mock_client = mock.MagicMock()524 with mock.patch.object(magics_impl._state, 'client', mock_client):525 magics_impl.grr_fgrep_impl(526 'foo bar', '/foo/bar', path_type=magics_impl.NTFS)527 mock_client.ntfs.fgrep.assert_called_once_with('/foo/bar', b'foo bar')528 def testRegistryPathType(self):529 mock_client = mock.MagicMock()530 with mock.patch.object(magics_impl._state, 'client', mock_client):531 magics_impl.grr_fgrep_impl(532 'foo bar', '/foo/bar', path_type=magics_impl.REGISTRY)533 mock_client.registry.fgrep.assert_called_once_with('/foo/bar', b'foo bar')534 def testHexString(self):535 mock_client = mock.MagicMock()536 with mock.patch.object(magics_impl._state, 'client', mock_client):537 magics_impl.grr_fgrep_impl('ffac90', '/foo/bar', hex_string=True)538 mock_client.os.fgrep.assert_called_once_with('/foo/bar', b'\xff\xac\x90')539class GrrInterrogateImplTest(absltest.TestCase):540 def testNoClientSelected(self):541 with self.assertRaises(magics_impl.NoClientSelectedError):542 magics_impl.grr_interrogate_impl()543 def testWithClientSet(self):544 mock_client = mock.MagicMock()545 with mock.patch.object(magics_impl._state, 'client', mock_client):546 magics_impl.grr_interrogate_impl()547 mock_client.interrogate.assert_called_once_with()548class GrrHostnameImplTest(absltest.TestCase):549 def testNoClientSelected(self):550 with self.assertRaises(magics_impl.NoClientSelectedError):551 magics_impl.grr_hostname_impl()552 def testWithClientSet(self):553 mock_client = _MockClient()554 mock_client.set_hostname('foobar')555 with mock.patch.object(magics_impl._state, 'client', mock_client):556 hostname = magics_impl.grr_hostname_impl()557 self.assertEqual(hostname, 'foobar')558class GrrIfconfigImplTest(absltest.TestCase):559 def testNoClientSelected(self):560 with self.assertRaises(magics_impl.NoClientSelectedError):561 magics_impl.grr_ifconfig_impl()562 def testWithClientSet(self):563 mock_client = _MockClient(ifaces=['foo', 'bar'])564 with mock.patch.object(magics_impl._state, 'client', mock_client):565 df = magics_impl.grr_ifconfig_impl()566 self.assertEqual(df.shape, (2, 1))567 self.assertEqual(list(df['ifname']), ['foo', 'bar'])568 def testPrettyIpAddress(self):569 mock_client = _MockClient()570 ipv4 = ipaddress.IPv4Address('42.0.255.32')571 ipv6 = ipaddress.IPv6Address('2001:db8::1000')572 mock_client.add_iface('foo', [ipv4.packed, ipv6.packed])573 mock_client.add_iface('bar', [])574 with mock.patch.object(magics_impl._state, 'client', mock_client):575 df = magics_impl.grr_ifconfig_impl()576 self.assertEqual(df.shape, (2, 2))577 self.assertEqual(list(df['ifname']), ['foo', 'bar'])578 self.assertTrue(pd.isna(df['addresses'][1]))579 addresses = df['addresses'][0]580 self.assertEqual(581 list(addresses['packed_bytes']), [ipv4.packed, ipv6.packed])582 self.assertEqual(583 list(addresses['packed_bytes.pretty']),584 [str(ipv4), str(ipv6)])585 def testPrettyMacAddress(self):586 mock_client = _MockClient()587 mock_client.add_iface('foo', mac=b'\xaa\x12\x42\xff\xa5\xd0')588 mock_client.add_iface('bar')589 with mock.patch.object(magics_impl._state, 'client', mock_client):590 df = magics_impl.grr_ifconfig_impl()591 self.assertEqual(df.shape, (2, 3))592 self.assertEqual(list(df['ifname']), ['foo', 'bar'])593 self.assertTrue(pd.isna(df['mac_address'][1]))594 self.assertTrue(pd.isna(df['mac_address.pretty'][1]))595 self.assertEqual(df['mac_address'][0], b'\xaa\x12\x42\xff\xa5\xd0')596 self.assertEqual(df['mac_address.pretty'][0], 'aa:12:42:ff:a5:d0')597class GrrUnameImplTest(absltest.TestCase):598 def testNoClientSelected(self):599 with self.assertRaises(magics_impl.NoClientSelectedError):600 magics_impl.grr_uname_impl()601 def testKernelRelease(self):602 mock_client = _MockClient(kernel='foobar')603 with mock.patch.object(magics_impl._state, 'client', mock_client):604 kernel_release = magics_impl.grr_uname_impl(kernel_release=True)605 self.assertEqual(kernel_release, 'foobar')606 def testMachine(self):607 mock_client = _MockClient(arch='foobar')608 with mock.patch.object(magics_impl._state, 'client', mock_client):609 machine = magics_impl.grr_uname_impl(machine=True)610 self.assertEqual(machine, 'foobar')611 def testNoOptionsProvided(self):612 mock_client = _MockClient()613 with mock.patch.object(magics_impl._state, 'client', mock_client):614 with self.assertRaises(ValueError):615 magics_impl.grr_uname_impl()616class GrrPsImplTest(absltest.TestCase):617 def testNoClientSelected(self):618 with self.assertRaises(magics_impl.NoClientSelectedError):619 magics_impl.grr_ps_impl()620 def testWithClientSet(self):621 mock_client = mock.MagicMock()622 with mock.patch.object(magics_impl._state, 'client', mock_client):623 magics_impl.grr_ps_impl()624 mock_client.ps.assert_called_once_with()625class GrrOsqueryiImplTest(absltest.TestCase):626 def testNoClientSelected(self):627 with self.assertRaises(magics_impl.NoClientSelectedError):628 magics_impl.grr_osqueryi_impl('foo bar')629 def testWithClientSet(self):630 mock_client = mock.MagicMock()631 with mock.patch.object(magics_impl._state, 'client', mock_client):632 magics_impl.grr_osqueryi_impl('foo bar')633 mock_client.osquery.assert_called_once_with('foo bar')634class GrrCollectImplTest(absltest.TestCase):635 def testNoClientSelected(self):636 with self.assertRaises(magics_impl.NoClientSelectedError):637 magics_impl.grr_collect_impl('FakeArtifact')638 def testWithClientSet(self):639 mock_client = mock.MagicMock()640 with mock.patch.object(magics_impl._state, 'client', mock_client):641 magics_impl.grr_collect_impl('FakeArtifact')642 mock_client.collect.assert_called_once_with('FakeArtifact')643class GrrYaraImplTest(absltest.TestCase):644 def testNoClientSelected(self):645 with self.assertRaises(magics_impl.NoClientSelectedError):646 magics_impl.grr_yara_impl('foo')647 def testWithClientSet(self):648 mock_client = mock.MagicMock()649 with mock.patch.object(magics_impl._state, 'client', mock_client):650 magics_impl.grr_yara_impl('foo', [42], 'bar')651 mock_client.yara.assert_called_once_with('foo', [42], 'bar')652class GrrWgetImplTest(absltest.TestCase):653 def testNoClientSelected(self):654 with self.assertRaises(magics_impl.NoClientSelectedError):655 magics_impl.grr_wget_impl('foo')656 def testRelativePath(self):657 mock_client = mock.MagicMock()658 with mock.patch.object(magics_impl._state, 'client', mock_client):659 with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):660 magics_impl.grr_wget_impl('foo/bar')661 mock_client.os.wget.assert_called_once_with('/quux/foo/bar')662 def testAbsolutePath(self):663 mock_client = mock.MagicMock()664 with mock.patch.object(magics_impl._state, 'client', mock_client):665 with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):666 magics_impl.grr_wget_impl('/foo/bar')667 mock_client.os.wget.assert_called_once_with('/foo/bar')668 def testCached(self):669 mock_client = mock.MagicMock()670 with mock.patch.object(magics_impl._state, 'client', mock_client):671 with mock.patch.object(magics_impl._state, 'cur_dir', '/quux'):672 magics_impl.grr_wget_impl('/foo/bar', cached=True)673 mock_client.os.cached.wget.assert_called_once_with('/foo/bar')674 def testTskPathType(self):675 mock_client = mock.MagicMock()676 with mock.patch.object(magics_impl._state, 'client', mock_client):677 magics_impl.grr_wget_impl('/foo/bar', path_type=magics_impl.TSK)678 mock_client.tsk.wget.assert_called_once_with('/foo/bar')679 def testNtfsPathType(self):680 mock_client = mock.MagicMock()681 with mock.patch.object(magics_impl._state, 'client', mock_client):682 magics_impl.grr_wget_impl('/foo/bar', path_type=magics_impl.NTFS)683 mock_client.ntfs.wget.assert_called_once_with('/foo/bar')684 def testRegistryPathType(self):685 mock_client = mock.MagicMock()686 with mock.patch.object(magics_impl._state, 'client', mock_client):687 magics_impl.grr_wget_impl('/foo/bar', path_type=magics_impl.REGISTRY)688 mock_client.registry.wget.assert_called_once_with('/foo/bar')689class _MockClient(grr_colab.Client):690 class MockInnerClient(object):691 def __init__(self):692 self.data = client_pb2.ApiClient()693 class MockVFS(object):694 def __init__(self):695 self.files = collections.defaultdict(bytes)696 def open(self, path):697 data = self.files[path]698 return io.BytesIO(data)699 def add_file(self, path, data):700 self.files[path] = data701 class MockFileSystem(object):702 def __init__(self):703 self._cached = _MockClient.MockVFS()704 self.files = collections.defaultdict(bytes)705 @property706 def cached(self):707 return self._cached708 def open(self, path):709 data = self.files[path]710 return io.BytesIO(data)711 def add_file(self, path, data):712 self.files[path] = data713 def __init__(self,714 client_id='',715 last_seen_at=0,716 ifaces=None,717 kernel='',718 arch=''):719 if ifaces is None:720 ifaces = []721 self._client = _MockClient.MockInnerClient()722 self._summary = None723 self._os = _MockClient.MockFileSystem()724 self._tsk = _MockClient.MockFileSystem()725 self._ntfs = _MockClient.MockFileSystem()726 self._registry = _MockClient.MockFileSystem()727 self.set_id(client_id)728 self.set_last_seen_at(last_seen_at)729 self.set_ifaces(ifaces)730 self.set_kernel(kernel)731 self.set_arch(arch)732 @property733 def os(self):734 return self._os735 @property736 def tsk(self):737 return self._tsk738 @property739 def ntfs(self):740 return self._ntfs741 @property742 def registry(self):743 return self._registry744 def set_id(self, client_id):745 self._client.data.client_id = client_id746 self._client.client_id = client_id747 def set_last_seen_at(self, ms):748 self._client.data.last_seen_at = ms749 def set_hostname(self, hostname):750 self._client.data.knowledge_base.fqdn = hostname751 def set_os_version(self, os_version):752 self._client.data.os_info.version = os_version753 def set_kernel(self, kernel):754 self._client.data.os_info.kernel = kernel755 def set_arch(self, arch):756 self._client.data.os_info.machine = arch757 def set_ifaces(self, ifnames):758 for ifname in ifnames:759 self.add_iface(ifname, [])760 def add_iface(self, ifname, addresses_bytes=None, mac=None):761 if addresses_bytes is None:762 addresses_bytes = []763 iface = self._client.data.interfaces.add()764 iface.ifname = ifname765 if mac is not None:766 iface.mac_address = mac767 for packed_bytes in addresses_bytes:768 addr = iface.addresses.add()769 addr.packed_bytes = packed_bytes770if __name__ == '__main__':...
test_influxdb.py
Source:test_influxdb.py
1"""The tests for the InfluxDB component."""2import datetime3import unittest4from unittest import mock5import influxdb as influx_client6from homeassistant.setup import setup_component7import homeassistant.components.influxdb as influxdb8from homeassistant.const import EVENT_STATE_CHANGED, STATE_OFF, STATE_ON, \9 STATE_STANDBY10from tests.common import get_test_home_assistant11@mock.patch('influxdb.InfluxDBClient')12@mock.patch(13 'homeassistant.components.influxdb.InfluxThread.batch_timeout',14 mock.Mock(return_value=0))15class TestInfluxDB(unittest.TestCase):16 """Test the InfluxDB component."""17 def setUp(self):18 """Set up things to be run when tests are started."""19 self.hass = get_test_home_assistant()20 self.handler_method = None21 self.hass.bus.listen = mock.Mock()22 def tearDown(self):23 """Clear data."""24 self.hass.stop()25 def test_setup_config_full(self, mock_client):26 """Test the setup with full configuration."""27 config = {28 'influxdb': {29 'host': 'host',30 'port': 123,31 'database': 'db',32 'username': 'user',33 'password': 'password',34 'max_retries': 4,35 'ssl': 'False',36 'verify_ssl': 'False',37 }38 }39 assert setup_component(self.hass, influxdb.DOMAIN, config)40 self.assertTrue(self.hass.bus.listen.called)41 self.assertEqual(42 EVENT_STATE_CHANGED, self.hass.bus.listen.call_args_list[0][0][0])43 self.assertTrue(mock_client.return_value.query.called)44 def test_setup_config_defaults(self, mock_client):45 """Test the setup with default configuration."""46 config = {47 'influxdb': {48 'host': 'host',49 'username': 'user',50 'password': 'pass',51 }52 }53 assert setup_component(self.hass, influxdb.DOMAIN, config)54 self.assertTrue(self.hass.bus.listen.called)55 self.assertEqual(56 EVENT_STATE_CHANGED, self.hass.bus.listen.call_args_list[0][0][0])57 def test_setup_minimal_config(self, mock_client):58 """Test the setup with minimal configuration."""59 config = {60 'influxdb': {}61 }62 assert setup_component(self.hass, influxdb.DOMAIN, config)63 def test_setup_missing_password(self, mock_client):64 """Test the setup with existing username and missing password."""65 config = {66 'influxdb': {67 'username': 'user'68 }69 }70 assert not setup_component(self.hass, influxdb.DOMAIN, config)71 def test_setup_query_fail(self, mock_client):72 """Test the setup for query failures."""73 config = {74 'influxdb': {75 'host': 'host',76 'username': 'user',77 'password': 'pass',78 }79 }80 mock_client.return_value.query.side_effect = \81 influx_client.exceptions.InfluxDBClientError('fake')82 assert not setup_component(self.hass, influxdb.DOMAIN, config)83 def _setup(self, **kwargs):84 """Set up the client."""85 config = {86 'influxdb': {87 'host': 'host',88 'username': 'user',89 'password': 'pass',90 'exclude': {91 'entities': ['fake.blacklisted'],92 'domains': ['another_fake']93 }94 }95 }96 config['influxdb'].update(kwargs)97 assert setup_component(self.hass, influxdb.DOMAIN, config)98 self.handler_method = self.hass.bus.listen.call_args_list[0][0][1]99 def test_event_listener(self, mock_client):100 """Test the event listener."""101 self._setup()102 # map of HA State to valid influxdb [state, value] fields103 valid = {104 '1': [None, 1],105 '1.0': [None, 1.0],106 STATE_ON: [STATE_ON, 1],107 STATE_OFF: [STATE_OFF, 0],108 STATE_STANDBY: [STATE_STANDBY, None],109 'foo': ['foo', None]110 }111 for in_, out in valid.items():112 attrs = {113 'unit_of_measurement': 'foobars',114 'longitude': '1.1',115 'latitude': '2.2',116 'battery_level': '99%',117 'temperature': '20c',118 'last_seen': 'Last seen 23 minutes ago',119 'updated_at': datetime.datetime(2017, 1, 1, 0, 0),120 'multi_periods': '0.120.240.2023873'121 }122 state = mock.MagicMock(123 state=in_, domain='fake', entity_id='fake.entity-id',124 object_id='entity', attributes=attrs)125 event = mock.MagicMock(data={'new_state': state}, time_fired=12345)126 body = [{127 'measurement': 'foobars',128 'tags': {129 'domain': 'fake',130 'entity_id': 'entity',131 },132 'time': 12345,133 'fields': {134 'longitude': 1.1,135 'latitude': 2.2,136 'battery_level_str': '99%',137 'battery_level': 99.0,138 'temperature_str': '20c',139 'temperature': 20.0,140 'last_seen_str': 'Last seen 23 minutes ago',141 'last_seen': 23.0,142 'updated_at_str': '2017-01-01 00:00:00',143 'updated_at': 20170101000000,144 'multi_periods_str': '0.120.240.2023873'145 },146 }]147 if out[0] is not None:148 body[0]['fields']['state'] = out[0]149 if out[1] is not None:150 body[0]['fields']['value'] = out[1]151 self.handler_method(event)152 self.hass.data[influxdb.DOMAIN].block_till_done()153 self.assertEqual(154 mock_client.return_value.write_points.call_count, 1155 )156 self.assertEqual(157 mock_client.return_value.write_points.call_args,158 mock.call(body)159 )160 mock_client.return_value.write_points.reset_mock()161 def test_event_listener_no_units(self, mock_client):162 """Test the event listener for missing units."""163 self._setup()164 for unit in (None, ''):165 if unit:166 attrs = {'unit_of_measurement': unit}167 else:168 attrs = {}169 state = mock.MagicMock(170 state=1, domain='fake', entity_id='fake.entity-id',171 object_id='entity', attributes=attrs)172 event = mock.MagicMock(data={'new_state': state}, time_fired=12345)173 body = [{174 'measurement': 'fake.entity-id',175 'tags': {176 'domain': 'fake',177 'entity_id': 'entity',178 },179 'time': 12345,180 'fields': {181 'value': 1,182 },183 }]184 self.handler_method(event)185 self.hass.data[influxdb.DOMAIN].block_till_done()186 self.assertEqual(187 mock_client.return_value.write_points.call_count, 1188 )189 self.assertEqual(190 mock_client.return_value.write_points.call_args,191 mock.call(body)192 )193 mock_client.return_value.write_points.reset_mock()194 def test_event_listener_inf(self, mock_client):195 """Test the event listener for missing units."""196 self._setup()197 attrs = {'bignumstring': '9' * 999, 'nonumstring': 'nan'}198 state = mock.MagicMock(199 state=8, domain='fake', entity_id='fake.entity-id',200 object_id='entity', attributes=attrs)201 event = mock.MagicMock(data={'new_state': state}, time_fired=12345)202 body = [{203 'measurement': 'fake.entity-id',204 'tags': {205 'domain': 'fake',206 'entity_id': 'entity',207 },208 'time': 12345,209 'fields': {210 'value': 8,211 },212 }]213 self.handler_method(event)214 self.hass.data[influxdb.DOMAIN].block_till_done()215 self.assertEqual(216 mock_client.return_value.write_points.call_count, 1217 )218 self.assertEqual(219 mock_client.return_value.write_points.call_args,220 mock.call(body)221 )222 mock_client.return_value.write_points.reset_mock()223 def test_event_listener_states(self, mock_client):224 """Test the event listener against ignored states."""225 self._setup()226 for state_state in (1, 'unknown', '', 'unavailable'):227 state = mock.MagicMock(228 state=state_state, domain='fake', entity_id='fake.entity-id',229 object_id='entity', attributes={})230 event = mock.MagicMock(data={'new_state': state}, time_fired=12345)231 body = [{232 'measurement': 'fake.entity-id',233 'tags': {234 'domain': 'fake',235 'entity_id': 'entity',236 },237 'time': 12345,238 'fields': {239 'value': 1,240 },241 }]242 self.handler_method(event)243 self.hass.data[influxdb.DOMAIN].block_till_done()244 if state_state == 1:245 self.assertEqual(246 mock_client.return_value.write_points.call_count, 1247 )248 self.assertEqual(249 mock_client.return_value.write_points.call_args,250 mock.call(body)251 )252 else:253 self.assertFalse(mock_client.return_value.write_points.called)254 mock_client.return_value.write_points.reset_mock()255 def test_event_listener_blacklist(self, mock_client):256 """Test the event listener against a blacklist."""257 self._setup()258 for entity_id in ('ok', 'blacklisted'):259 state = mock.MagicMock(260 state=1, domain='fake', entity_id='fake.{}'.format(entity_id),261 object_id=entity_id, attributes={})262 event = mock.MagicMock(data={'new_state': state}, time_fired=12345)263 body = [{264 'measurement': 'fake.{}'.format(entity_id),265 'tags': {266 'domain': 'fake',267 'entity_id': entity_id,268 },269 'time': 12345,270 'fields': {271 'value': 1,272 },273 }]274 self.handler_method(event)275 self.hass.data[influxdb.DOMAIN].block_till_done()276 if entity_id == 'ok':277 self.assertEqual(278 mock_client.return_value.write_points.call_count, 1279 )280 self.assertEqual(281 mock_client.return_value.write_points.call_args,282 mock.call(body)283 )284 else:285 self.assertFalse(mock_client.return_value.write_points.called)286 mock_client.return_value.write_points.reset_mock()287 def test_event_listener_blacklist_domain(self, mock_client):288 """Test the event listener against a blacklist."""289 self._setup()290 for domain in ('ok', 'another_fake'):291 state = mock.MagicMock(292 state=1, domain=domain,293 entity_id='{}.something'.format(domain),294 object_id='something', attributes={})295 event = mock.MagicMock(data={'new_state': state}, time_fired=12345)296 body = [{297 'measurement': '{}.something'.format(domain),298 'tags': {299 'domain': domain,300 'entity_id': 'something',301 },302 'time': 12345,303 'fields': {304 'value': 1,305 },306 }]307 self.handler_method(event)308 self.hass.data[influxdb.DOMAIN].block_till_done()309 if domain == 'ok':310 self.assertEqual(311 mock_client.return_value.write_points.call_count, 1312 )313 self.assertEqual(314 mock_client.return_value.write_points.call_args,315 mock.call(body)316 )317 else:318 self.assertFalse(mock_client.return_value.write_points.called)319 mock_client.return_value.write_points.reset_mock()320 def test_event_listener_whitelist(self, mock_client):321 """Test the event listener against a whitelist."""322 config = {323 'influxdb': {324 'host': 'host',325 'username': 'user',326 'password': 'pass',327 'include': {328 'entities': ['fake.included'],329 }330 }331 }332 assert setup_component(self.hass, influxdb.DOMAIN, config)333 self.handler_method = self.hass.bus.listen.call_args_list[0][0][1]334 for entity_id in ('included', 'default'):335 state = mock.MagicMock(336 state=1, domain='fake', entity_id='fake.{}'.format(entity_id),337 object_id=entity_id, attributes={})338 event = mock.MagicMock(data={'new_state': state}, time_fired=12345)339 body = [{340 'measurement': 'fake.{}'.format(entity_id),341 'tags': {342 'domain': 'fake',343 'entity_id': entity_id,344 },345 'time': 12345,346 'fields': {347 'value': 1,348 },349 }]350 self.handler_method(event)351 self.hass.data[influxdb.DOMAIN].block_till_done()352 if entity_id == 'included':353 self.assertEqual(354 mock_client.return_value.write_points.call_count, 1355 )356 self.assertEqual(357 mock_client.return_value.write_points.call_args,358 mock.call(body)359 )360 else:361 self.assertFalse(mock_client.return_value.write_points.called)362 mock_client.return_value.write_points.reset_mock()363 def test_event_listener_whitelist_domain(self, mock_client):364 """Test the event listener against a whitelist."""365 config = {366 'influxdb': {367 'host': 'host',368 'username': 'user',369 'password': 'pass',370 'include': {371 'domains': ['fake'],372 }373 }374 }375 assert setup_component(self.hass, influxdb.DOMAIN, config)376 self.handler_method = self.hass.bus.listen.call_args_list[0][0][1]377 for domain in ('fake', 'another_fake'):378 state = mock.MagicMock(379 state=1, domain=domain,380 entity_id='{}.something'.format(domain),381 object_id='something', attributes={})382 event = mock.MagicMock(data={'new_state': state}, time_fired=12345)383 body = [{384 'measurement': '{}.something'.format(domain),385 'tags': {386 'domain': domain,387 'entity_id': 'something',388 },389 'time': 12345,390 'fields': {391 'value': 1,392 },393 }]394 self.handler_method(event)395 self.hass.data[influxdb.DOMAIN].block_till_done()396 if domain == 'fake':397 self.assertEqual(398 mock_client.return_value.write_points.call_count, 1399 )400 self.assertEqual(401 mock_client.return_value.write_points.call_args,402 mock.call(body)403 )404 else:405 self.assertFalse(mock_client.return_value.write_points.called)406 mock_client.return_value.write_points.reset_mock()407 def test_event_listener_invalid_type(self, mock_client):408 """Test the event listener when an attribute has an invalid type."""409 self._setup()410 # map of HA State to valid influxdb [state, value] fields411 valid = {412 '1': [None, 1],413 '1.0': [None, 1.0],414 STATE_ON: [STATE_ON, 1],415 STATE_OFF: [STATE_OFF, 0],416 STATE_STANDBY: [STATE_STANDBY, None],417 'foo': ['foo', None]418 }419 for in_, out in valid.items():420 attrs = {421 'unit_of_measurement': 'foobars',422 'longitude': '1.1',423 'latitude': '2.2',424 'invalid_attribute': ['value1', 'value2']425 }426 state = mock.MagicMock(427 state=in_, domain='fake', entity_id='fake.entity-id',428 object_id='entity', attributes=attrs)429 event = mock.MagicMock(data={'new_state': state}, time_fired=12345)430 body = [{431 'measurement': 'foobars',432 'tags': {433 'domain': 'fake',434 'entity_id': 'entity',435 },436 'time': 12345,437 'fields': {438 'longitude': 1.1,439 'latitude': 2.2,440 'invalid_attribute_str': "['value1', 'value2']"441 },442 }]443 if out[0] is not None:444 body[0]['fields']['state'] = out[0]445 if out[1] is not None:446 body[0]['fields']['value'] = out[1]447 self.handler_method(event)448 self.hass.data[influxdb.DOMAIN].block_till_done()449 self.assertEqual(450 mock_client.return_value.write_points.call_count, 1451 )452 self.assertEqual(453 mock_client.return_value.write_points.call_args,454 mock.call(body)455 )456 mock_client.return_value.write_points.reset_mock()457 def test_event_listener_default_measurement(self, mock_client):458 """Test the event listener with a default measurement."""459 config = {460 'influxdb': {461 'host': 'host',462 'username': 'user',463 'password': 'pass',464 'default_measurement': 'state',465 'exclude': {466 'entities': ['fake.blacklisted']467 }468 }469 }470 assert setup_component(self.hass, influxdb.DOMAIN, config)471 self.handler_method = self.hass.bus.listen.call_args_list[0][0][1]472 for entity_id in ('ok', 'blacklisted'):473 state = mock.MagicMock(474 state=1, domain='fake', entity_id='fake.{}'.format(entity_id),475 object_id=entity_id, attributes={})476 event = mock.MagicMock(data={'new_state': state}, time_fired=12345)477 body = [{478 'measurement': 'state',479 'tags': {480 'domain': 'fake',481 'entity_id': entity_id,482 },483 'time': 12345,484 'fields': {485 'value': 1,486 },487 }]488 self.handler_method(event)489 self.hass.data[influxdb.DOMAIN].block_till_done()490 if entity_id == 'ok':491 self.assertEqual(492 mock_client.return_value.write_points.call_count, 1493 )494 self.assertEqual(495 mock_client.return_value.write_points.call_args,496 mock.call(body)497 )498 else:499 self.assertFalse(mock_client.return_value.write_points.called)500 mock_client.return_value.write_points.reset_mock()501 def test_event_listener_unit_of_measurement_field(self, mock_client):502 """Test the event listener for unit of measurement field."""503 config = {504 'influxdb': {505 'host': 'host',506 'username': 'user',507 'password': 'pass',508 'override_measurement': 'state',509 }510 }511 assert setup_component(self.hass, influxdb.DOMAIN, config)512 self.handler_method = self.hass.bus.listen.call_args_list[0][0][1]513 attrs = {514 'unit_of_measurement': 'foobars',515 }516 state = mock.MagicMock(517 state='foo', domain='fake', entity_id='fake.entity-id',518 object_id='entity', attributes=attrs)519 event = mock.MagicMock(data={'new_state': state}, time_fired=12345)520 body = [{521 'measurement': 'state',522 'tags': {523 'domain': 'fake',524 'entity_id': 'entity',525 },526 'time': 12345,527 'fields': {528 'state': 'foo',529 'unit_of_measurement_str': 'foobars',530 },531 }]532 self.handler_method(event)533 self.hass.data[influxdb.DOMAIN].block_till_done()534 self.assertEqual(535 mock_client.return_value.write_points.call_count, 1536 )537 self.assertEqual(538 mock_client.return_value.write_points.call_args,539 mock.call(body)540 )541 mock_client.return_value.write_points.reset_mock()542 def test_event_listener_tags_attributes(self, mock_client):543 """Test the event listener when some attributes should be tags."""544 config = {545 'influxdb': {546 'host': 'host',547 'username': 'user',548 'password': 'pass',549 'tags_attributes': ['friendly_fake']550 }551 }552 assert setup_component(self.hass, influxdb.DOMAIN, config)553 self.handler_method = self.hass.bus.listen.call_args_list[0][0][1]554 attrs = {555 'friendly_fake': 'tag_str',556 'field_fake': 'field_str',557 }558 state = mock.MagicMock(559 state=1, domain='fake',560 entity_id='fake.something',561 object_id='something', attributes=attrs)562 event = mock.MagicMock(data={'new_state': state}, time_fired=12345)563 body = [{564 'measurement': 'fake.something',565 'tags': {566 'domain': 'fake',567 'entity_id': 'something',568 'friendly_fake': 'tag_str'569 },570 'time': 12345,571 'fields': {572 'value': 1,573 'field_fake_str': 'field_str'574 },575 }]576 self.handler_method(event)577 self.hass.data[influxdb.DOMAIN].block_till_done()578 self.assertEqual(579 mock_client.return_value.write_points.call_count, 1580 )581 self.assertEqual(582 mock_client.return_value.write_points.call_args,583 mock.call(body)584 )585 mock_client.return_value.write_points.reset_mock()586 def test_event_listener_component_override_measurement(self, mock_client):587 """Test the event listener with overridden measurements."""588 config = {589 'influxdb': {590 'host': 'host',591 'username': 'user',592 'password': 'pass',593 'component_config': {594 'sensor.fake_humidity': {595 'override_measurement': 'humidity'596 }597 },598 'component_config_glob': {599 'binary_sensor.*motion': {600 'override_measurement': 'motion'601 }602 },603 'component_config_domain': {604 'climate': {605 'override_measurement': 'hvac'606 }607 }608 }609 }610 assert setup_component(self.hass, influxdb.DOMAIN, config)611 self.handler_method = self.hass.bus.listen.call_args_list[0][0][1]612 test_components = [613 {'domain': 'sensor', 'id': 'fake_humidity', 'res': 'humidity'},614 {'domain': 'binary_sensor', 'id': 'fake_motion', 'res': 'motion'},615 {'domain': 'climate', 'id': 'fake_thermostat', 'res': 'hvac'},616 {'domain': 'other', 'id': 'just_fake', 'res': 'other.just_fake'},617 ]618 for comp in test_components:619 state = mock.MagicMock(620 state=1, domain=comp['domain'],621 entity_id=comp['domain'] + '.' + comp['id'],622 object_id=comp['id'], attributes={})623 event = mock.MagicMock(data={'new_state': state}, time_fired=12345)624 body = [{625 'measurement': comp['res'],626 'tags': {627 'domain': comp['domain'],628 'entity_id': comp['id']629 },630 'time': 12345,631 'fields': {632 'value': 1,633 },634 }]635 self.handler_method(event)636 self.hass.data[influxdb.DOMAIN].block_till_done()637 self.assertEqual(638 mock_client.return_value.write_points.call_count, 1639 )640 self.assertEqual(641 mock_client.return_value.write_points.call_args,642 mock.call(body)643 )644 mock_client.return_value.write_points.reset_mock()645 def test_scheduled_write(self, mock_client):646 """Test the event listener to retry after write failures."""647 config = {648 'influxdb': {649 'host': 'host',650 'username': 'user',651 'password': 'pass',652 'max_retries': 1653 }654 }655 assert setup_component(self.hass, influxdb.DOMAIN, config)656 self.handler_method = self.hass.bus.listen.call_args_list[0][0][1]657 state = mock.MagicMock(658 state=1, domain='fake', entity_id='entity.id', object_id='entity',659 attributes={})660 event = mock.MagicMock(data={'new_state': state}, time_fired=12345)661 mock_client.return_value.write_points.side_effect = \662 IOError('foo')663 # Write fails664 with mock.patch.object(influxdb.time, 'sleep') as mock_sleep:665 self.handler_method(event)666 self.hass.data[influxdb.DOMAIN].block_till_done()667 assert mock_sleep.called668 json_data = mock_client.return_value.write_points.call_args[0][0]669 self.assertEqual(mock_client.return_value.write_points.call_count, 2)670 mock_client.return_value.write_points.assert_called_with(json_data)671 # Write works again672 mock_client.return_value.write_points.side_effect = None673 with mock.patch.object(influxdb.time, 'sleep') as mock_sleep:674 self.handler_method(event)675 self.hass.data[influxdb.DOMAIN].block_till_done()676 assert not mock_sleep.called677 self.assertEqual(mock_client.return_value.write_points.call_count, 3)678 def test_queue_backlog_full(self, mock_client):679 """Test the event listener to drop old events."""680 self._setup()681 state = mock.MagicMock(682 state=1, domain='fake', entity_id='entity.id', object_id='entity',683 attributes={})684 event = mock.MagicMock(data={'new_state': state}, time_fired=12345)685 monotonic_time = 0686 def fast_monotonic():687 """Monotonic time that ticks fast enough to cause a timeout."""688 nonlocal monotonic_time689 monotonic_time += 60690 return monotonic_time691 with mock.patch('homeassistant.components.influxdb.time.monotonic',692 new=fast_monotonic):693 self.handler_method(event)694 self.hass.data[influxdb.DOMAIN].block_till_done()695 self.assertEqual(696 mock_client.return_value.write_points.call_count, 0697 )...
benchmark_uploader_test.py
Source:benchmark_uploader_test.py
1# Copyright 2017 The TensorFlow Authors. All Rights Reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14# ==============================================================================15"""Tests for benchmark_uploader."""16from __future__ import absolute_import17from __future__ import division18from __future__ import print_function19import json20import os21import tempfile22import unittest23from mock import MagicMock24from mock import patch25import tensorflow as tf # pylint: disable=g-bad-import-order26try:27 from google.cloud import bigquery28 from official.benchmark import benchmark_uploader29except ImportError:30 bigquery = None31 benchmark_uploader = None32@unittest.skipIf(bigquery is None, "Bigquery dependency is not installed.")33class BigQueryUploaderTest(tf.test.TestCase):34 @patch.object(bigquery, "Client")35 def setUp(self, mock_bigquery):36 self.mock_client = mock_bigquery.return_value37 self.mock_dataset = MagicMock(name="dataset")38 self.mock_table = MagicMock(name="table")39 self.mock_client.dataset.return_value = self.mock_dataset40 self.mock_dataset.table.return_value = self.mock_table41 self.mock_client.insert_rows_json.return_value = []42 self.benchmark_uploader = benchmark_uploader.BigQueryUploader()43 self.benchmark_uploader._bq_client = self.mock_client44 self.log_dir = tempfile.mkdtemp(dir=self.get_temp_dir())45 with open(os.path.join(self.log_dir, "metric.log"), "a") as f:46 json.dump({"name": "accuracy", "value": 1.0}, f)47 f.write("\n")48 json.dump({"name": "loss", "value": 0.5}, f)49 f.write("\n")50 with open(os.path.join(self.log_dir, "run.log"), "w") as f:51 json.dump({"model_name": "value"}, f)52 def tearDown(self):53 tf.gfile.DeleteRecursively(self.get_temp_dir())54 def test_upload_benchmark_run_json(self):55 self.benchmark_uploader.upload_benchmark_run_json(56 "dataset", "table", "run_id", {"model_name": "value"})57 self.mock_client.insert_rows_json.assert_called_once_with(58 self.mock_table, [{"model_name": "value", "model_id": "run_id"}])59 def test_upload_benchmark_metric_json(self):60 metric_json_list = [61 {"name": "accuracy", "value": 1.0},62 {"name": "loss", "value": 0.5}63 ]64 expected_params = [65 {"run_id": "run_id", "name": "accuracy", "value": 1.0},66 {"run_id": "run_id", "name": "loss", "value": 0.5}67 ]68 self.benchmark_uploader.upload_benchmark_metric_json(69 "dataset", "table", "run_id", metric_json_list)70 self.mock_client.insert_rows_json.assert_called_once_with(71 self.mock_table, expected_params)72 def test_upload_benchmark_run_file(self):73 self.benchmark_uploader.upload_benchmark_run_file(74 "dataset", "table", "run_id", os.path.join(self.log_dir, "run.log"))75 self.mock_client.insert_rows_json.assert_called_once_with(76 self.mock_table, [{"model_name": "value", "model_id": "run_id"}])77 def test_upload_metric_file(self):78 self.benchmark_uploader.upload_metric_file(79 "dataset", "table", "run_id",80 os.path.join(self.log_dir, "metric.log"))81 expected_params = [82 {"run_id": "run_id", "name": "accuracy", "value": 1.0},83 {"run_id": "run_id", "name": "loss", "value": 0.5}84 ]85 self.mock_client.insert_rows_json.assert_called_once_with(86 self.mock_table, expected_params)87 def test_insert_run_status(self):88 self.benchmark_uploader.insert_run_status(89 "dataset", "table", "run_id", "status")90 expected_query = ("INSERT dataset.table "91 "(run_id, status) "92 "VALUES('run_id', 'status')")93 self.mock_client.query.assert_called_once_with(query=expected_query)94 def test_update_run_status(self):95 self.benchmark_uploader.update_run_status(96 "dataset", "table", "run_id", "status")97 expected_query = ("UPDATE dataset.table "98 "SET status = 'status' "99 "WHERE run_id = 'run_id'")100 self.mock_client.query.assert_called_once_with(query=expected_query)101if __name__ == "__main__":...
Automation Testing Tutorials
Learn to execute automation testing from scratch with LambdaTest Learning Hub. Right from setting up the prerequisites to run your first automation test, to following best practices and diving deeper into advanced test scenarios. LambdaTest Learning Hubs compile a list of step-by-step guides to help you be proficient with different test automation frameworks i.e. Selenium, Cypress, TestNG etc.
LambdaTest Learning Hubs:
- JUnit Tutorial
- TestNG Tutorial
- Webdriver Tutorial
- WebDriverIO Tutorial
- Protractor Tutorial
- Selenium 4 Tutorial
- Jenkins Tutorial
- NUnit Tutorial
- Jest Tutorial
- Playwright Tutorial
- Cypress Tutorial
- PyTest Tutorial
YouTube
You could also refer to video tutorials over LambdaTest YouTube channel to get step by step demonstration from industry experts.
Run tempest automation tests on LambdaTest cloud grid
Perform automation testing on 3000+ real desktop and mobile devices online.
Try LambdaTest Now !!
Get 100 minutes of automation test minutes FREE!!
Was this article helpful?
