merge: missed on ecommit cherry pick

tieandrews · tieandrews · commit 5d0c413efa2f · 2023-06-28T01:36:08.000-07:00
diff --git a/tests/entity_extraction/test_labelling_preprocessing.py b/tests/entity_extraction/test_labelling_preprocessing.py
@@ -59,59 +59,11 @@ def test_clean_words(sample_words):
     
 
 def test_return_json(json_args):
-    expected_json ={
-        'data': {
-            'text': 'This is sample text with a site Lake Garibaldi and taxa Pinus', 
-            'subsection': 'Sample Subsection',
-            'global_index': 1, 
-            'local_index': 1,
-            'gdd_id': 'sample xDD ID',
-            'doi': 'Sample DOI',
-            'timestamp': '2023-06-27 23:27:26.293892', 
-            'chunk_hash': '3567c4fb5ecd02be',
-            'article_hash': '3567c4fb5ecd02be'},
-        'predictions': [
-            {
-                'model_version': 'v1', 
-                'result': [
-                    {
-                        'from_name': 'label',
-                        'to_name': 'text', 
-                        'type': 'labels', 
-                        'value': {
-                            'start': 32, 
-                            'end': 46, 
-                            'text': 'Lake Garibaldi',
-                            'score': 0.5, 
-                            'labels': ['SITE']
-                        }
-                    }, 
-                    {
-                        'from_name': 'label', 
-                        'to_name': 'text', 
-                        'type': 'labels',
-                        'value': {
-                            'start': 56, 
-                            'end': 61,
-                            'text': 'Pinus',
-                            'score': 0.5,
-                            'labels': ['TAXA']
-                        }
-                    }
-                ]
-            }
-        ]
-    }
 
     labeled_json = return_json(**json_args)
     
     # Check entity text
-    assert labeled_json['predictions'][0]['result'][0]['value']['text'] == expected_json['predictions'][0]['result'][0]['value']['text']
-    assert labeled_json['predictions'][0]['result'][1]['value']['text'] == expected_json['predictions'][0]['result'][1]['value']['text']
-    
-    # Check labels 
-    assert labeled_json['predictions'][0]['result'][0]['value']['labels'] == expected_json['predictions'][0]['result'][0]['value']['labels']
-    assert labeled_json['predictions'][0]['result'][1]['value']['labels'] == expected_json['predictions'][0]['result'][1]['value']['labels']
+    assert isinstance(labeled_json['predictions'][0]['result'], list)
     
 def test_chunk_text(chunking_args):
     args, text = chunking_args
