|
21 | 21 | import time |
22 | 22 | import subprocess |
23 | 23 | import shutil |
| 24 | +from hashlib import sha256 |
24 | 25 |
|
25 | 26 | from gevent import socket, sleep, spawn |
26 | 27 |
|
@@ -353,3 +354,83 @@ def test_sftp_mkdir_rel_path(self): |
353 | 354 | shutil.rmtree(_dir) |
354 | 355 | except Exception: |
355 | 356 | pass |
| 357 | + |
| 358 | + def test_scp_recv_large_file(self): |
| 359 | + cur_dir = os.path.dirname(__file__) |
| 360 | + file_name = 'file1' |
| 361 | + file_copy_to = 'file_copied' |
| 362 | + file_path_from = os.path.sep.join([cur_dir, file_name]) |
| 363 | + file_copy_to_dirpath = os.path.expanduser('~/') + file_copy_to |
| 364 | + for _path in (file_path_from, file_copy_to_dirpath): |
| 365 | + try: |
| 366 | + os.unlink(_path) |
| 367 | + except OSError: |
| 368 | + pass |
| 369 | + try: |
| 370 | + with open(file_path_from, 'wb') as fh: |
| 371 | + # ~300MB |
| 372 | + for _ in range(20000000): |
| 373 | + fh.write(b"adsfasldkfjabafj") |
| 374 | + self.client.scp_recv(file_path_from, file_copy_to_dirpath) |
| 375 | + self.assertTrue(os.path.isfile(file_copy_to_dirpath)) |
| 376 | + read_file_size = os.stat(file_path_from).st_size |
| 377 | + written_file_size = os.stat(file_copy_to_dirpath).st_size |
| 378 | + self.assertEqual(read_file_size, written_file_size) |
| 379 | + sha = sha256() |
| 380 | + with open(file_path_from, 'rb') as fh: |
| 381 | + for block in fh: |
| 382 | + sha.update(block) |
| 383 | + read_file_hash = sha.hexdigest() |
| 384 | + sha = sha256() |
| 385 | + with open(file_copy_to_dirpath, 'rb') as fh: |
| 386 | + for block in fh: |
| 387 | + sha.update(block) |
| 388 | + written_file_hash = sha.hexdigest() |
| 389 | + self.assertEqual(read_file_hash, written_file_hash) |
| 390 | + finally: |
| 391 | + for _path in (file_path_from, file_copy_to_dirpath): |
| 392 | + try: |
| 393 | + os.unlink(_path) |
| 394 | + except Exception: |
| 395 | + pass |
| 396 | + |
| 397 | + def test_scp_send_large_file(self): |
| 398 | + cur_dir = os.path.dirname(__file__) |
| 399 | + file_name = 'file1' |
| 400 | + file_copy_to = 'file_copied' |
| 401 | + file_path_from = os.path.sep.join([cur_dir, file_name]) |
| 402 | + file_copy_to_dirpath = os.path.expanduser('~/') + file_copy_to |
| 403 | + for _path in (file_path_from, file_copy_to_dirpath): |
| 404 | + try: |
| 405 | + os.unlink(_path) |
| 406 | + except OSError: |
| 407 | + pass |
| 408 | + try: |
| 409 | + with open(file_path_from, 'wb') as fh: |
| 410 | + # ~300MB |
| 411 | + for _ in range(20000000): |
| 412 | + fh.write(b"adsfasldkfjabafj") |
| 413 | + self.client.scp_send(file_path_from, file_copy_to_dirpath) |
| 414 | + self.assertTrue(os.path.isfile(file_copy_to_dirpath)) |
| 415 | + # OS file flush race condition |
| 416 | + sleep(.1) |
| 417 | + read_file_size = os.stat(file_path_from).st_size |
| 418 | + written_file_size = os.stat(file_copy_to_dirpath).st_size |
| 419 | + self.assertEqual(read_file_size, written_file_size) |
| 420 | + sha = sha256() |
| 421 | + with open(file_path_from, 'rb') as fh: |
| 422 | + for block in fh: |
| 423 | + sha.update(block) |
| 424 | + read_file_hash = sha.hexdigest() |
| 425 | + sha = sha256() |
| 426 | + with open(file_copy_to_dirpath, 'rb') as fh: |
| 427 | + for block in fh: |
| 428 | + sha.update(block) |
| 429 | + written_file_hash = sha.hexdigest() |
| 430 | + self.assertEqual(read_file_hash, written_file_hash) |
| 431 | + finally: |
| 432 | + for _path in (file_path_from, file_copy_to_dirpath): |
| 433 | + try: |
| 434 | + os.unlink(_path) |
| 435 | + except Exception: |
| 436 | + pass |
0 commit comments