-
Notifications
You must be signed in to change notification settings - Fork 7
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Signed-off-by: Ryan Koo <[email protected]>
- Loading branch information
Showing
10 changed files
with
543 additions
and
478 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,23 @@ | ||
# | ||
# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. | ||
# | ||
|
||
import os | ||
import unittest | ||
|
||
from aistore import Client | ||
|
||
from utils import generate_random_str | ||
|
||
class TestBase(unittest.TestCase): | ||
def setUp(self): | ||
self.endpoint = os.environ.get("AIS_ENDPOINT", "http://192.168.49.2:8080") | ||
self.git_test_mode = os.getenv('GIT_TEST', 'False') | ||
self.client = Client(self.endpoint) | ||
self.test_bck = self.client.bucket("etl-test-bucket").create(exist_ok=True) | ||
self.test_etl = self.client.etl("test-etl-" + generate_random_str()) | ||
|
||
def tearDown(self): | ||
self.test_bck.delete() | ||
self.test_etl.stop() | ||
self.test_etl.delete() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,186 @@ | ||
# | ||
# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. | ||
# | ||
|
||
import bz2 | ||
import gzip | ||
import hashlib | ||
import os | ||
import unittest | ||
|
||
from aistore.sdk.etl_const import ETL_COMM_HPULL | ||
from aistore.sdk.etl_templates import COMPRESS | ||
|
||
from test_base import TestBase | ||
from utils import git_test_mode_format_image_tag_test | ||
|
||
class TestCompressTransformer(TestBase): | ||
def setUp(self): | ||
super().setUp() | ||
self.test_image_filename = "test-image.jpg" | ||
self.test_image_source = "./resources/test-image.jpg" | ||
self.test_text_filename = "test-text.txt" | ||
self.test_text_source = "./resources/test-text.txt" | ||
self.test_image_gz_filename = "test-image.jpg.gz" | ||
self.test_image_gz_source = "./resources/test-image.jpg.gz" | ||
self.test_text_gz_filename = "test-text.txt.gz" | ||
self.test_text_gz_source = "./resources/test-text.txt.gz" | ||
self.test_image_bz2_filename = "test-image.jpg.bz2" | ||
self.test_image_bz2_source = "./resources/test-image.jpg.bz2" | ||
self.test_text_bz2_filename = "test-text.txt.bz2" | ||
self.test_text_bz2_source = "./resources/test-text.txt.bz2" | ||
self.test_text_bz2_filename = "test-text.txt.bz2" | ||
self.test_text_bz2_source = "./resources/test-text.txt.bz2" | ||
self.test_bck.object(self.test_image_filename).put_file(self.test_image_source) | ||
self.test_bck.object(self.test_text_filename).put_file(self.test_text_source) | ||
self.test_bck.object(self.test_image_gz_filename).put_file(self.test_image_gz_source) | ||
self.test_bck.object(self.test_text_gz_filename).put_file(self.test_text_gz_source) | ||
self.test_bck.object(self.test_image_bz2_filename).put_file(self.test_image_bz2_source) | ||
self.test_bck.object(self.test_text_bz2_filename).put_file(self.test_text_bz2_source) | ||
|
||
def tearDown(self): | ||
super().tearDown() | ||
|
||
@unittest.skipIf(os.getenv('COMPRESS_ENABLE', 'true') == 'false', "COMPRESS is disabled") | ||
def test_compress_gzip(self): | ||
template = COMPRESS.format(communication_type=ETL_COMM_HPULL, arg1="--mode", val1="compress", arg2="--compression", val2="gzip") | ||
|
||
if self.git_test_mode == 'true': | ||
template = git_test_mode_format_image_tag_test(template, "compress") | ||
|
||
self.test_etl.init_spec(template=template, communication_type=ETL_COMM_HPULL) | ||
|
||
compressed_image = self.test_bck.object(self.test_image_filename).get(etl_name=self.test_etl.name).read_all() | ||
compressed_text = self.test_bck.object(self.test_text_filename).get(etl_name=self.test_etl.name).read_all() | ||
|
||
self.assertNotEqual(compressed_image, b"Data processing failed") | ||
self.assertNotEqual(compressed_text, b"Data processing failed") | ||
|
||
# Decompress the files | ||
decompressed_image = gzip.decompress(compressed_image) | ||
decompressed_text = gzip.decompress(compressed_text) | ||
|
||
with open(self.test_image_source, 'rb') as file: | ||
original_image_content = file.read() | ||
|
||
with open(self.test_text_source, 'r') as file: | ||
original_text_content = file.read() | ||
|
||
self.assertEqual(decompressed_image, original_image_content) | ||
self.assertEqual(decompressed_text.decode('utf-8'), original_text_content) | ||
|
||
# Calculate the checksums | ||
original_image_checksum = hashlib.md5(original_image_content).hexdigest() | ||
decompressed_image_checksum = hashlib.md5(decompressed_image).hexdigest() | ||
original_text_checksum = hashlib.md5(original_text_content.encode('utf-8')).hexdigest() | ||
decompressed_text_checksum = hashlib.md5(decompressed_text).hexdigest() | ||
|
||
# Validate the checksums | ||
self.assertEqual(original_image_checksum, decompressed_image_checksum) | ||
self.assertEqual(original_text_checksum, decompressed_text_checksum) | ||
|
||
@unittest.skipIf(os.getenv('COMPRESS_ENABLE', 'true') == 'false', "COMPRESS is disabled") | ||
def test_compress_bz2(self): | ||
template = COMPRESS.format(communication_type=ETL_COMM_HPULL, arg1="--mode", val1="compress", arg2="--compression", val2="bz2") | ||
|
||
if self.git_test_mode == 'true': | ||
template = git_test_mode_format_image_tag_test(template, "compress") | ||
|
||
self.test_etl.init_spec(template=template, communication_type=ETL_COMM_HPULL) | ||
|
||
compressed_image = self.test_bck.object(self.test_image_filename).get(etl_name=self.test_etl.name).read_all() | ||
compressed_text = self.test_bck.object(self.test_text_filename).get(etl_name=self.test_etl.name).read_all() | ||
|
||
self.assertNotEqual(compressed_image, b"Data processing failed") | ||
self.assertNotEqual(compressed_text, b"Data processing failed") | ||
|
||
# Decompress the files | ||
decompressed_image = bz2.decompress(compressed_image) | ||
decompressed_text = bz2.decompress(compressed_text) | ||
|
||
with open(self.test_image_source, 'rb') as file: | ||
original_image_content = file.read() | ||
|
||
with open(self.test_text_source, 'r') as file: | ||
original_text_content = file.read() | ||
|
||
self.assertEqual(decompressed_image, original_image_content) | ||
self.assertEqual(decompressed_text.decode('utf-8'), original_text_content) | ||
|
||
# Calculate the checksums | ||
original_image_checksum = hashlib.md5(original_image_content).hexdigest() | ||
decompressed_image_checksum = hashlib.md5(decompressed_image).hexdigest() | ||
original_text_checksum = hashlib.md5(original_text_content.encode('utf-8')).hexdigest() | ||
decompressed_text_checksum = hashlib.md5(decompressed_text).hexdigest() | ||
|
||
# Validate the checksums | ||
self.assertEqual(original_image_checksum, decompressed_image_checksum) | ||
self.assertEqual(original_text_checksum, decompressed_text_checksum) | ||
|
||
@unittest.skipIf(os.getenv('COMPRESS_ENABLE', 'true') == 'false', "COMPRESS is disabled") | ||
def test_decompress_gzip(self): | ||
template = COMPRESS.format(communication_type=ETL_COMM_HPULL, arg1="--mode", val1="decompress", arg2="--compression", val2="gzip") | ||
|
||
if self.git_test_mode == 'true': | ||
template = git_test_mode_format_image_tag_test(template, "compress") | ||
|
||
self.test_etl.init_spec(template=template, communication_type=ETL_COMM_HPULL) | ||
|
||
decompressed_image = self.test_bck.object(self.test_image_gz_filename).get(etl_name=self.test_etl.name).read_all() | ||
decompressed_text = self.test_bck.object(self.test_text_gz_filename).get(etl_name=self.test_etl.name).read_all() | ||
|
||
self.assertNotEqual(decompressed_image, b"Data processing failed") | ||
self.assertNotEqual(decompressed_text, b"Data processing failed") | ||
|
||
with open(self.test_image_source, 'rb') as file: | ||
original_image_content = file.read() | ||
|
||
with open(self.test_text_source, 'r') as file: | ||
original_text_content = file.read() | ||
|
||
self.assertEqual(decompressed_image, original_image_content) | ||
self.assertEqual(decompressed_text.decode('utf-8'), original_text_content) | ||
|
||
# Calculate the checksums | ||
original_image_checksum = hashlib.md5(original_image_content).hexdigest() | ||
decompressed_image_checksum = hashlib.md5(decompressed_image).hexdigest() | ||
original_text_checksum = hashlib.md5(original_text_content.encode('utf-8')).hexdigest() | ||
decompressed_text_checksum = hashlib.md5(decompressed_text).hexdigest() | ||
|
||
# Validate the checksums | ||
self.assertEqual(original_image_checksum, decompressed_image_checksum) | ||
self.assertEqual(original_text_checksum, decompressed_text_checksum) | ||
|
||
@unittest.skipIf(os.getenv('COMPRESS_ENABLE', 'true') == 'false', "COMPRESS is disabled") | ||
def test_decompress_bz2(self): | ||
template = COMPRESS.format(communication_type=ETL_COMM_HPULL, arg1="--mode", val1="decompress", arg2="--compression", val2="bz2") | ||
|
||
if self.git_test_mode == 'true': | ||
template = git_test_mode_format_image_tag_test(template, "compress") | ||
|
||
self.test_etl.init_spec(template=template, communication_type=ETL_COMM_HPULL) | ||
|
||
decompressed_image = self.test_bck.object(self.test_image_bz2_filename).get(etl_name=self.test_etl.name).read_all() | ||
decompressed_text = self.test_bck.object(self.test_text_bz2_filename).get(etl_name=self.test_etl.name).read_all() | ||
|
||
self.assertNotEqual(decompressed_image, b"Data processing failed") | ||
self.assertNotEqual(decompressed_text, b"Data processing failed") | ||
|
||
with open(self.test_image_source, 'rb') as file: | ||
original_image_content = file.read() | ||
|
||
with open(self.test_text_source, 'r') as file: | ||
original_text_content = file.read() | ||
|
||
self.assertEqual(decompressed_image, original_image_content) | ||
self.assertEqual(decompressed_text.decode('utf-8'), original_text_content) | ||
|
||
# Calculate the checksums | ||
original_image_checksum = hashlib.md5(original_image_content).hexdigest() | ||
decompressed_image_checksum = hashlib.md5(decompressed_image).hexdigest() | ||
original_text_checksum = hashlib.md5(original_text_content.encode('utf-8')).hexdigest() | ||
decompressed_text_checksum = hashlib.md5(decompressed_text).hexdigest() | ||
|
||
# Validate the checksums | ||
self.assertEqual(original_image_checksum, decompressed_image_checksum) | ||
self.assertEqual(original_text_checksum, decompressed_text_checksum) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,47 @@ | ||
# | ||
# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. | ||
# | ||
|
||
import os | ||
import unittest | ||
|
||
from aistore.sdk.etl_const import ETL_COMM_HPULL | ||
from aistore.sdk.etl_templates import ECHO | ||
|
||
from test_base import TestBase | ||
from utils import git_test_mode_format_image_tag_test | ||
|
||
class TestEchoTransformer(TestBase): | ||
def setUp(self): | ||
super().setUp() | ||
self.test_image_filename = "test-image.jpg" | ||
self.test_image_source = "./resources/test-image.jpg" | ||
self.test_text_filename = "test-text.txt" | ||
self.test_text_source = "./resources/test-text.txt" | ||
self.test_bck.object(self.test_image_filename).put_file(self.test_image_source) | ||
self.test_bck.object(self.test_text_filename).put_file(self.test_text_source) | ||
|
||
def tearDown(self): | ||
super().tearDown() | ||
|
||
@unittest.skipIf(os.getenv('ECHO_ENABLE', 'true') == 'false', "ECHO is disabled") | ||
def test_echo(self): | ||
template = ECHO.format(communication_type=ETL_COMM_HPULL) | ||
|
||
if self.git_test_mode == 'true': | ||
template = git_test_mode_format_image_tag_test(template, "echo") | ||
|
||
self.test_etl.init_spec(template=template, communication_type=ETL_COMM_HPULL) | ||
|
||
transformed_image_bytes = self.test_bck.object(self.test_image_filename).get(etl_name=self.test_etl.name).read_all() | ||
transformed_text_bytes = self.test_bck.object(self.test_text_filename).get(etl_name=self.test_etl.name).read_all() | ||
|
||
# Compare image content | ||
with open(self.test_image_source, 'rb') as file: | ||
original_image_content = file.read() | ||
self.assertEqual(transformed_image_bytes, original_image_content) | ||
|
||
# Compare text content | ||
with open(self.test_text_source, 'r') as file: | ||
original_text_content = file.read() | ||
self.assertEqual(transformed_text_bytes.decode('utf-8'), original_text_content) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,47 @@ | ||
# | ||
# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. | ||
# | ||
|
||
import os | ||
import unittest | ||
|
||
from aistore.sdk.etl_const import ETL_COMM_HPULL | ||
from aistore.sdk.etl_templates import GO_ECHO | ||
|
||
from test_base import TestBase | ||
from utils import git_test_mode_format_image_tag_test | ||
|
||
class TestGoEchoTransformer(TestBase): | ||
def setUp(self): | ||
super().setUp() | ||
self.test_image_filename = "test-image.jpg" | ||
self.test_image_source = "./resources/test-image.jpg" | ||
self.test_text_filename = "test-text.txt" | ||
self.test_text_source = "./resources/test-text.txt" | ||
self.test_bck.object(self.test_image_filename).put_file(self.test_image_source) | ||
self.test_bck.object(self.test_text_filename).put_file(self.test_text_source) | ||
|
||
def tearDown(self): | ||
super().tearDown() | ||
|
||
@unittest.skipIf(os.getenv('GO_ECHO_ENABLE', 'true') == 'false', "GO_ECHO is disabled") | ||
def test_go_echo(self): | ||
template = GO_ECHO.format(communication_type=ETL_COMM_HPULL) | ||
|
||
if self.git_test_mode == 'true': | ||
template = git_test_mode_format_image_tag_test(template, "echo_go") | ||
|
||
self.test_etl.init_spec(template=template, communication_type=ETL_COMM_HPULL) | ||
|
||
transformed_image_bytes = self.test_bck.object(self.test_image_filename).get(etl_name=self.test_etl.name).read_all() | ||
transformed_text_bytes = self.test_bck.object(self.test_text_filename).get(etl_name=self.test_etl.name).read_all() | ||
|
||
# Compare image content | ||
with open(self.test_image_source, 'rb') as file: | ||
original_image_content = file.read() | ||
self.assertEqual(transformed_image_bytes, original_image_content) | ||
|
||
# Compare text content | ||
with open(self.test_text_source, 'r') as file: | ||
original_text_content = file.read() | ||
self.assertEqual(transformed_text_bytes.decode('utf-8'), original_text_content) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,41 @@ | ||
# | ||
# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. | ||
# | ||
|
||
import os | ||
import unittest | ||
|
||
from aistore.sdk.etl_const import ETL_COMM_HPULL | ||
from aistore.sdk.etl_templates import HELLO_WORLD | ||
|
||
from test_base import TestBase | ||
from utils import git_test_mode_format_image_tag_test | ||
|
||
class TestHelloWorldTransformer(TestBase): | ||
def setUp(self): | ||
super().setUp() | ||
self.test_image_filename = "test-image.jpg" | ||
self.test_image_source = "./resources/test-image.jpg" | ||
self.test_text_filename = "test-text.txt" | ||
self.test_text_source = "./resources/test-text.txt" | ||
self.test_bck.object(self.test_image_filename).put_file(self.test_image_source) | ||
self.test_bck.object(self.test_text_filename).put_file(self.test_text_source) | ||
|
||
def tearDown(self): | ||
super().tearDown() | ||
|
||
@unittest.skipIf(os.getenv('HELLO_WORLD_ENABLE', 'true') == 'false', "HELLO_WORLD is disabled") | ||
def test_hello_world(self): | ||
template = HELLO_WORLD.format(communication_type=ETL_COMM_HPULL) | ||
|
||
if self.git_test_mode == 'true': | ||
template = git_test_mode_format_image_tag_test(template, "hello_world") | ||
|
||
self.test_etl.init_spec(template=template, communication_type=ETL_COMM_HPULL) | ||
|
||
transformed_image_bytes = self.test_bck.object(self.test_image_filename).get(etl_name=self.test_etl.name).read_all() | ||
transformed_text_bytes = self.test_bck.object(self.test_text_filename).get(etl_name=self.test_etl.name).read_all() | ||
|
||
# Compare file contents | ||
self.assertEqual(b"Hello World!", transformed_image_bytes) | ||
self.assertEqual(b"Hello World!", transformed_text_bytes) |
Oops, something went wrong.