diff --git a/tom_dataproducts/processors/spectroscopy_processor.py b/tom_dataproducts/processors/spectroscopy_processor.py index 410d6a03b..1b1e9314b 100644 --- a/tom_dataproducts/processors/spectroscopy_processor.py +++ b/tom_dataproducts/processors/spectroscopy_processor.py @@ -34,15 +34,15 @@ def process_data(self, data_product): mimetype = mimetypes.guess_type(data_product.data.path)[0] if mimetype in self.FITS_MIMETYPES: - spectrum, obs_date = self._process_spectrum_from_fits(data_product) + spectrum, obs_date, source_id = self._process_spectrum_from_fits(data_product) elif mimetype in self.PLAINTEXT_MIMETYPES: - spectrum, obs_date = self._process_spectrum_from_plaintext(data_product) + spectrum, obs_date, source_id = self._process_spectrum_from_plaintext(data_product) else: raise InvalidFileFormatException('Unsupported file type') serialized_spectrum = SpectrumSerializer().serialize(spectrum) - return [(obs_date, serialized_spectrum)] + return [(obs_date, serialized_spectrum, source_id)] def _process_spectrum_from_fits(self, data_product): """ @@ -61,12 +61,14 @@ def _process_spectrum_from_fits(self, data_product): datetime otherwise :rtype: AstroPy.Time """ + facility_name = 'DEFAULT' flux, header = fits.getdata(data_product.data.path, header=True) for facility_class in get_service_classes(): facility = get_service_class(facility_class)() if facility.is_fits_facility(header): + facility_name = facility_class flux_constant = facility.get_flux_constant() date_obs = facility.get_date_obs_from_fits_header(header) break @@ -86,7 +88,7 @@ def _process_spectrum_from_fits(self, data_product): spectrum = Spectrum1D(flux=flux, wcs=wcs) - return spectrum, Time(date_obs).to_datetime() + return spectrum, Time(date_obs).to_datetime(), facility_name def _process_spectrum_from_plaintext(self, data_product): """ @@ -132,4 +134,4 @@ def _process_spectrum_from_plaintext(self, data_product): flux = np.array(data['flux']) * flux_constant spectrum = Spectrum1D(flux=flux, spectral_axis=spectral_axis) - return spectrum, Time(date_obs).to_datetime() + return spectrum, Time(date_obs).to_datetime(), facility_name diff --git a/tom_dataproducts/tests/tests.py b/tom_dataproducts/tests/tests.py index 95a987dcc..d53d62fd6 100644 --- a/tom_dataproducts/tests/tests.py +++ b/tom_dataproducts/tests/tests.py @@ -447,7 +447,7 @@ def setUp(self): self.test_file = SimpleUploadedFile('afile.fits', b'somedata') @patch('tom_dataproducts.processors.spectroscopy_processor.SpectroscopyProcessor._process_spectrum_from_fits', - return_value=('', '')) + return_value=('', '', '')) @patch('tom_dataproducts.processors.spectroscopy_processor.SpectrumSerializer.serialize', return_value={}) def test_process_spectroscopy_with_fits_file(self, serializer_mock, process_data_mock): self.data_product.data.save('spectrum.fits', self.test_file) @@ -455,7 +455,7 @@ def test_process_spectroscopy_with_fits_file(self, serializer_mock, process_data process_data_mock.assert_called_with(self.data_product) @patch('tom_dataproducts.processors.spectroscopy_processor.SpectroscopyProcessor._process_spectrum_from_plaintext', - return_value=('', '')) + return_value=('', '', '')) @patch('tom_dataproducts.processors.spectroscopy_processor.SpectrumSerializer.serialize', return_value={}) def test_process_spectroscopy_with_plaintext_file(self, serializer_mock, process_data_mock): self.data_product.data.save('spectrum.csv', self.test_file) @@ -470,7 +470,7 @@ def test_process_spectroscopy_with_invalid_file_type(self): def test_process_spectrum_from_fits(self): with open('tom_dataproducts/tests/test_data/test_spectrum.fits', 'rb') as spectrum_file: self.data_product.data.save('spectrum.fits', spectrum_file) - spectrum, _ = self.spectrum_data_processor._process_spectrum_from_fits(self.data_product) + spectrum, _, _ = self.spectrum_data_processor._process_spectrum_from_fits(self.data_product) self.assertTrue(isinstance(spectrum, Spectrum1D)) self.assertAlmostEqual(spectrum.flux.mean().value, 2.295068e-14, places=19) self.assertAlmostEqual(spectrum.wavelength.mean().value, 6600.478789, places=5) @@ -478,7 +478,7 @@ def test_process_spectrum_from_fits(self): def test_process_spectrum_from_plaintext(self): with open('tom_dataproducts/tests/test_data/test_spectrum.csv', 'rb') as spectrum_file: self.data_product.data.save('spectrum.csv', spectrum_file) - spectrum, _ = self.spectrum_data_processor._process_spectrum_from_plaintext(self.data_product) + spectrum, _, _ = self.spectrum_data_processor._process_spectrum_from_plaintext(self.data_product) self.assertTrue(isinstance(spectrum, Spectrum1D)) self.assertAlmostEqual(spectrum.flux.mean().value, 1.166619e-14, places=19) self.assertAlmostEqual(spectrum.wavelength.mean().value, 3250.744489, places=5)