From 56089a734df17dd6547bab796a9c3a51cb99f7ea Mon Sep 17 00:00:00 2001
From: Noa Aviel Dove <nadove@ucsc.edu>
Date: Wed, 10 Apr 2024 01:58:54 -0700
Subject: [PATCH] Refactor canned file loading

---
 test/indexer/__init__.py |  8 ++++++--
 test/pfb_test_case.py    | 17 ++++++-----------
 2 files changed, 12 insertions(+), 13 deletions(-)

diff --git a/test/indexer/__init__.py b/test/indexer/__init__.py
index b7e199c1f9..d406a99cf1 100644
--- a/test/indexer/__init__.py
+++ b/test/indexer/__init__.py
@@ -6,6 +6,7 @@
 import os
 from typing import (
     Generic,
+    Literal,
     Optional,
     Type,
     Union,
@@ -89,6 +90,10 @@ class CannedFileTestCase(AzulUnitTestCase):
     expected outputs.
     """
 
+    @classmethod
+    def _data_path(cls, module: Literal['service', 'indexer'], path: str) -> str | bytes:
+        return os.path.join(config.project_root, 'test', module, 'data', path)
+
     @classmethod
     def _load_canned_file(cls,
                           bundle: BundleFQID,
@@ -111,10 +116,9 @@ def _load_canned_file_version(cls,
                                   version: Optional[str],
                                   extension: str
                                   ) -> Union[MutableJSONs, MutableJSON]:
-        data_prefix = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'data')
         suffix = '' if version is None else '.' + version
         file_name = f'{uuid}{suffix}.{extension}.json'
-        with open(os.path.join(data_prefix, file_name), 'r') as infile:
+        with open(cls._data_path('indexer', file_name), 'r') as infile:
             return json.load(infile)
 
 
diff --git a/test/pfb_test_case.py b/test/pfb_test_case.py
index 262d7d2d41..1fae557153 100644
--- a/test/pfb_test_case.py
+++ b/test/pfb_test_case.py
@@ -1,17 +1,14 @@
 import json
-from pathlib import (
-    Path,
-)
-import sys
+import os
 
 import fastavro
 
-from azul_test_case import (
-    AzulUnitTestCase,
+from indexer import (
+    CannedFileTestCase,
 )
 
 
-class PFBTestCase(AzulUnitTestCase):
+class PFBTestCase(CannedFileTestCase):
 
     def _assert_pfb_schema(self, schema):
         fastavro.parse_schema(schema)
@@ -22,10 +19,8 @@ def _assert_pfb_schema(self, schema):
         def to_json(records):
             return json.dumps(records, indent=4, sort_keys=True)
 
-        cls = type(self)
-        module = sys.modules[cls.__module__]
-        results_file = Path(module.__file__).parent / 'data' / 'pfb_manifest.schema.json'
-        if results_file.exists():
+        results_file = self._data_path('service', 'pfb_manifest.schema.json')
+        if os.path.exists(results_file):
             with open(results_file, 'r') as f:
                 expected_records = json.load(f)
             self.assertEqual(expected_records, json.loads(to_json(schema)))