diff --git a/src/parsing/asp/grammar.pest b/src/parsing/asp/grammar.pest
index cf266244..b347088c 100644
--- a/src/parsing/asp/grammar.pest
+++ b/src/parsing/asp/grammar.pest
@@ -51,4 +51,5 @@ body = { (atomic_formula ~ (("," | ";") ~ atomic_formula)*)? }
 
 rule = { (!"." ~ head ~ (":-" ~ body)?) ~ "." }
 
-program = { rule* }
+// TODO: Think about a proper way to handle the whitespace at the end of programs
+program = { rule* ~ WHITESPACE* }
diff --git a/src/parsing/asp/pest.rs b/src/parsing/asp/pest.rs
index 5516f7ad..85614444 100644
--- a/src/parsing/asp/pest.rs
+++ b/src/parsing/asp/pest.rs
@@ -1026,6 +1026,18 @@ mod tests {
                     ],
                 },
             ),
+            (
+                "a.\n",
+                Program {
+                    rules: vec![Rule {
+                        head: Head::Basic(Atom {
+                            predicate: "a".into(),
+                            terms: vec![],
+                        }),
+                        body: Body { formulas: vec![] },
+                    }],
+                },
+            ),
         ]);
     }
 }
diff --git a/src/parsing/fol/grammar.pest b/src/parsing/fol/grammar.pest
index 6e6f34d1..0b54794a 100644
--- a/src/parsing/fol/grammar.pest
+++ b/src/parsing/fol/grammar.pest
@@ -8,7 +8,7 @@
 
 WHITESPACE = _{ " " | NEWLINE }
 
-keyword = _{ primitive | binary_connective | unary_connective | quantifier | sort }
+keyword = _{ primitive | binary_connective | unary_connective | quantifier }
 
     primitive = _{ infimum | supremum }
         infimum  = { "#inf" }
@@ -79,4 +79,5 @@ formula = { prefix* ~ primary ~ (infix ~ prefix* ~ primary)* }
     infix   = _{ binary_connective }
     primary = _{ "(" ~ formula ~ ")" | atomic_formula }
 
-theory = { (formula ~ ".")+ }
+// TODO: Think about a proper way to handle the whitespace at the end of programs
+theory = { (formula ~ ".")* ~ WHITESPACE* }
diff --git a/src/parsing/fol/pest.rs b/src/parsing/fol/pest.rs
index 6f277694..4b5514f2 100644
--- a/src/parsing/fol/pest.rs
+++ b/src/parsing/fol/pest.rs
@@ -480,14 +480,14 @@ mod tests {
             AtomParser, AtomicFormulaParser, BasicIntegerTermParser, BinaryConnectiveParser,
             BinaryOperatorParser, ComparisonParser, FormulaParser, GeneralTermParser, GuardParser,
             IntegerTermParser, QuantificationParser, QuantifierParser, RelationParser,
-            UnaryConnectiveParser, UnaryOperatorParser, VariableParser,
+            TheoryParser, UnaryConnectiveParser, UnaryOperatorParser, VariableParser,
         },
         crate::{
             parsing::TestedParser,
             syntax_tree::fol::{
                 Atom, AtomicFormula, BasicIntegerTerm, BinaryConnective, BinaryOperator,
                 Comparison, Formula, GeneralTerm, Guard, IntegerTerm, Quantification, Quantifier,
-                Relation, Sort, UnaryConnective, UnaryOperator, Variable,
+                Relation, Sort, Theory, UnaryConnective, UnaryOperator, Variable,
             },
         },
     };
@@ -756,6 +756,14 @@ mod tests {
                         terms: vec![],
                     },
                 ),
+                (
+                    // Parsing "g" caused issues ealier because "g" is also a sort declaration.
+                    "g",
+                    Atom {
+                        predicate: "g".into(),
+                        terms: vec![],
+                    },
+                ),
                 (
                     "p()",
                     Atom {
@@ -1139,7 +1147,21 @@ mod tests {
                 },
             ),
         ]);
+    }
 
-        // TODO Add more tests
+    #[test]
+    fn parse_theory() {
+        TheoryParser.should_parse_into([
+            ("", Theory { formulas: vec![] }),
+            (
+                "a.\n",
+                Theory {
+                    formulas: vec![Formula::AtomicFormula(AtomicFormula::Atom(Atom {
+                        predicate: "a".into(),
+                        terms: vec![],
+                    }))],
+                },
+            ),
+        ]);
     }
 }